import gradio as gr
from diffusers import StableDiffusionPipeline
import torch

# Detect device
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float16 if device == "cuda" else torch.float32

# Load base model
pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=dtype
)

# Load LoRA weights
pipe.load_lora_weights("./lora")

# Move to device
pipe = pipe.to(device)

# Extra CPU optimizations
if device == "cpu":
    pipe.enable_attention_slicing()
    pipe.enable_sequential_cpu_offload()
    pipe.enable_vae_tiling()

# Define image generation function with quality toggle
def generate(prompt, quality):
    if quality == "Fast":
        steps = 20
        guidance_scale = 7.0
    else:  # High Quality
        steps = 40
        guidance_scale = 8.5
    
    with torch.no_grad():
        image = pipe(prompt, num_inference_steps=steps, guidance_scale=guidance_scale).images[0]
    return image

# Build Gradio UI
demo = gr.Interface(
    fn=generate,
    inputs=[
        gr.Textbox(label="Enter your prompt"),
        gr.Dropdown(["Fast", "High Quality"], value="Fast", label="Generation Mode")
    ],
    outputs=gr.Image(label="Generated Image"),
    title="Fine-tuned Stable Diffusion with LoRA (CPU-Optimized)",
    description="Choose 'Fast' for quicker generation or 'High Quality' for better details."
)

# Launch app
if __name__ == "__main__":
    demo.launch()