import torch
import random
import gradio as gr
from diffusers import StableDiffusionPipeline, DDIMScheduler

# Fine-tuned model weights path
model_path = "model"

# Load fine-tuned pipeline
pipe = StableDiffusionPipeline.from_pretrained(model_path, torch_dtype=torch.float32)

# Change Scheduler to DDIMScheduler
pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)

# For memory optimization -- only for GPU
# pipe.enable_xformers_memory_efficient_attention()

# Generate the image from the given text prompt
def generate_response(prompt):

    # define the negative prompt
    negative_prompt = "bad anatomy, ugly, deformed, desfigured, distorted face, poorly drawn, blurry, low quality, low definition, lowres, out of frame, out of image, cropped, cut off, signature, watermark"

    # initialize the variables
    num_samples = 1
    guidance_scale = 7.5
    num_inference_steps = 10
    height = 512
    width = 512
    
    seed = random.randint(0, 2147483647)
    generator = torch.Generator(device='cpu').manual_seed(seed)
    
    with torch.inference_mode():
        imgs = pipe(
            prompt,
            negative_prompt=negative_prompt,
            height=height, 
            width=width,
            num_images_per_prompt=num_samples,
            num_inference_steps=num_inference_steps,
            guidance_scale=guidance_scale,
            generator=generator
        ).images

    return imgs[0]

# Create and launch the Gradio UI    
gradio_ui = gr.Interface(fn=generate_response, inputs="text", outputs="image")
gradio_ui.launch()