import gradio as gr
from diffusers import StableDiffusionPipeline, EulerDiscreteScheduler
import torch
from datetime import datetime

# Check for GPU and set up the model
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float16 if device == "cuda" else torch.float32

# Use a smaller, faster model
model_id = "runwayml/stable-diffusion-v1-5"

# Load the pipeline with optimizations
scheduler = EulerDiscreteScheduler.from_pretrained(model_id, subfolder="scheduler")
pipe = StableDiffusionPipeline.from_pretrained(
    model_id,
    scheduler=scheduler,
    torch_dtype=dtype,
    safety_checker=None,  # Disable safety checker for faster generation
    requires_safety_checker=False
).to(device)

# Apply optimizations based on available hardware
if device == "cuda":
    try:
        pipe.enable_xformers_memory_efficient_attention()  # Flash attention if available
    except:
        print("xformers not available, using default attention")
    pipe.enable_attention_slicing(1)  # Minimal slicing for balance between speed and memory
else:
    pipe.enable_attention_slicing()  # More aggressive slicing for CPU

# Warm up the model (important for consistent speed)
print("Warming up the model...")
with torch.inference_mode():
    _ = pipe("warmup", num_inference_steps=1, guidance_scale=1)

def generate_image(
    prompt: str,
    negative_prompt: str = "",
    steps: int = 20,  # Slightly more steps for CPU
    guidance_scale: float = 7.0,
    seed: int = -1,
    width: int = 512,
    height: int = 512
):
    # Start timer
    start_time = datetime.now()
    
    # Set random seed if provided
    generator = None
    if seed != -1:
        generator = torch.Generator(device=device).manual_seed(seed)
    
    # Generate image with optimizations
    with torch.inference_mode(), torch.autocast(device) if device == "cuda" else torch.no_grad():
        image = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            num_inference_steps=int(steps),
            guidance_scale=guidance_scale,
            generator=generator,
            width=width,
            height=height
        ).images[0]
    
    # Calculate generation time
    gen_time = (datetime.now() - start_time).total_seconds()
    print(f"Generated image in {gen_time:.2f} seconds")
    
    return image, f"Device: {device.upper()} | Time: {gen_time:.2f}s | Steps: {steps} | Size: {width}x{height}"

# Gradio interface
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # ⚡ Diffusion Image Generator
    *Optimized for Hugging Face Spaces*
    """)
    
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(
                label="Prompt",
                placeholder="A beautiful landscape with mountains and a lake",
                max_lines=3
            )
            negative_prompt = gr.Textbox(
                label="Negative Prompt",
                placeholder="blurry, low quality, distorted",
                max_lines=2
            )
            
            with gr.Accordion("Advanced Settings", open=False):
                with gr.Row():
                    steps = gr.Slider(8, 30, value=20, step=1, label="Inference Steps")
                    guidance_scale = gr.Slider(1, 10, value=7.0, step=0.5, label="Guidance Scale")
                with gr.Row():
                    seed = gr.Number(value=-1, label="Seed (-1 for random)")
                    width = gr.Slider(384, 512, value=512, step=64, label="Width")
                    height = gr.Slider(384, 512, value=512, step=64, label="Height")
            
            generate_btn = gr.Button("Generate Image", variant="primary")
            info_output = gr.Textbox(label="Generation Info", interactive=False)
        
        with gr.Column():
            output_image = gr.Image(label="Generated Image", type="pil", show_download_button=True)
    
    # Examples for quick testing
    gr.Examples(
        examples=[
            ["A futuristic cyberpunk city at night, neon lights, rain reflections", "blurry, low quality", 20, 7.0, 42],
            ["A cute robot cat, digital art, vibrant colors", "ugly, deformed", 18, 7.5, -1],
            ["A majestic dragon flying over mountains at sunset, fantasy art", "cartoon, sketch", 25, 8.0, 123]
        ],
        inputs=[prompt, negative_prompt, steps, guidance_scale, seed],
        outputs=[output_image, info_output],
        fn=generate_image,
        cache_examples=True
    )
    
    generate_btn.click(
        fn=generate_image,
        inputs=[prompt, negative_prompt, steps, guidance_scale, seed, width, height],
        outputs=[output_image, info_output]
    )

# For Hugging Face Spaces
demo.queue(max_size=2)  # More conservative queue size
demo.launch(debug=False)