import gradio as gr
import torch
from diffusers import StableDiffusionPipeline

# Base model
model_id = "runwayml/stable-diffusion-v1-5"
pipe = StableDiffusionPipeline.from_pretrained(
    model_id,
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
)

# Load LoRA
pipe.load_lora_weights("yashu16/pokemon-lora-v1")

# Move to GPU if available
if torch.cuda.is_available():
    pipe = pipe.to("cuda")

def generate(prompt, steps, guidance, size, seed):
    if seed == -1 or seed is None:
        generator = None
    else:
        generator = torch.Generator("cuda" if torch.cuda.is_available() else "cpu").manual_seed(int(seed))

    image = pipe(
        prompt,
        num_inference_steps=int(steps),
        guidance_scale=float(guidance),
        height=int(size),
        width=int(size),
        generator=generator
    ).images[0]
    return image

with gr.Blocks() as demo:
    gr.Markdown("## 🎨 Pokemon LoRA Generator")
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(label="📝 Prompt", value="generate a Pikachu Pokemon")
            steps = gr.Slider(1, 50, value=20, step=1, label="Inference Steps")
            guidance = gr.Slider(1, 15, value=7.5, step=0.1, label="Guidance Scale")
            size = gr.Radio([256, 512, 768], value=512, label="Image Size")
            seed = gr.Number(value=42, label="Seed (-1 for random)")
            btn = gr.Button("🚀 Generate")
        with gr.Column():
            output = gr.Image(label="Generated Pokemon")

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

demo.launch()