import gradio as gr
from model import sd_model  # Import the model wrapper
import torch
from datetime import datetime

def generate_image(prompt, negative_prompt="", steps=25, guidance=7.5, seed=None):
    try:
        generator = torch.Generator("cuda").manual_seed(int(seed)) if seed else None
        return sd_model.generate(
            prompt,
            negative_prompt=negative_prompt,
            num_inference_steps=int(steps),
            guidance_scale=float(guidance),
            generator=generator
        )
    except Exception as e:
        raise gr.Error(f"Generation failed: {str(e)}")

with gr.Blocks() as app:
    gr.Markdown("# 🖼️ Stable Diffusion Image Generator")
    
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(label="Prompt")
            negative_prompt = gr.Textbox(label="Negative Prompt")
            steps = gr.Slider(10, 50, value=25)
            guidance = gr.Slider(1, 20, value=7.5)
            seed = gr.Number(label="Seed (blank for random)")
            btn = gr.Button("Generate")
        
        with gr.Column():
            output = gr.Image()

    btn.click(
        fn=generate_image,
        inputs=[prompt, negative_prompt, steps, guidance, seed],
        outputs=output
    )

if __name__ == "__main__":
    app.launch()