import os
import torch
import gradio as gr
import numpy as np

from PIL import Image

from diffusers import StableDiffusionPipeline,UNet2DConditionModel

NEGATIVE_PROMPT = "worst quality, low quality, bad anatomy, watermark, text, blurry, cartoon, unreal"

unet = UNet2DConditionModel.from_pretrained("runwayml/stable-diffusion-v1-5",subfolder='unet').to("cuda")



# unet.load_lora_weights("./exp_output/celeba_finetune/checkpoint-20000", weight_name="pytorch_lora_weights.safetensors")

pipeline = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    unet=unet)

pipeline.load_lora_weights("./exp_output/celeba_finetune/checkpoint-20000", weight_name="pytorch_lora_weights.safetensors")

# Define a function to process input and return output
def generate_image(text,num_batch,is_use_lora,num_inference_steps):
    # Process text to generate image
    if is_use_lora:
        pipeline.enable_lora()
    else:
        pipeline.disable_lora()
    
    print('begin inference with text:', text, 'is_use_lora:', is_use_lora)
    image = pipeline(text,
                     num_inference_steps=num_inference_steps,
                     num_images_per_prompt=num_batch,
                     negative_prompt=NEGATIVE_PROMPT).images
    return image


with gr.Blocks() as demo:
    
    with gr.Row():
        with gr.Column():
            with gr.Row():
                is_use_lora = gr.Checkbox(label="Use LoRA", value=False)
                num_batch = gr.Number(value=4,label="Number of batch")
                num_inference_steps = gr.Number(value=20,label="Number of inference steps")
    
            text_input = gr.Textbox(lines=2, label="Input text", value="A young woman with long hair and a big smile.")
            generate_button = gr.Button(value="Generate image")

        # image_out = gr.Image(label="Output image", height=512,width=512)
        image_out = gr.Gallery(label="Generated images", show_label=False, elem_id="gallery", object_fit="contain", height="512")
    
    generate_button.click(generate_image, inputs=[text_input,num_batch,is_use_lora,num_inference_steps], outputs=image_out)

demo.launch(server_port=7861)