import gradio as gr
from diffusers import StableDiffusionPipeline
import torch

# Configuration
MODEL_ID = "runwayml/stable-diffusion-v1-5"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32

print(f"Loading model on {DEVICE} with {DTYPE}...")

# Load pipeline with optimizations
# safety_checker=None saves memory and storage
pipe = StableDiffusionPipeline.from_pretrained(
    MODEL_ID,
    torch_dtype=DTYPE,
    safety_checker=None,
    requires_safety_checker=False,
    use_safetensors=True
)

pipe.to(DEVICE)

# Enable memory optimizations
if DEVICE == "cuda":
    pipe.enable_xformers_memory_efficient_attention()
else:
    pipe.enable_attention_slicing()

def generate(prompt, seed):
    if not prompt: 
        return None
    
    # Handle seed
    generator = None
    if seed != -1:
        generator = torch.Generator(DEVICE).manual_seed(int(seed))
    
    # Generate
    image = pipe(
        prompt, 
        width=512, 
        height=512,
        num_inference_steps=25,
        generator=generator
    ).images[0]
    
    return image

# Minimal Interface
with gr.Blocks(title="SD v1.5 Lite") as demo:
    gr.Markdown("# Fast & Lite Stable Diffusion 🚀")
    
    with gr.Row():
        with gr.Column():
            txt_prompt = gr.Textbox(label="Prompt", placeholder="A cyberpunk cat...", lines=2)
            num_seed = gr.Number(label="Seed (-1 for random)", value=-1, precision=0)
            btn_run = gr.Button("Generate", variant="primary")
        
        with gr.Column():
            img_out = gr.Image(label="Result")
            
    btn_run.click(fn=generate, inputs=[txt_prompt, num_seed], outputs=img_out)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)