import gradio as gr
import torch
from diffusers import StableDiffusionPipeline

# Load model
print("Loading model...")
pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    safety_checker=None
)

# Load YOUR LoRA weights
pipe.load_lora_weights("ozzyzoz123/indian-clothing-lora")

# Move to GPU if available
device = "cuda" if torch.cuda.is_available() else "cpu"
pipe = pipe.to(device)
print(f"Model loaded on {device}")

# Generation function
def generate(prompt, negative_prompt, steps, guidance_scale):
    image = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        num_inference_steps=int(steps),
        guidance_scale=guidance_scale
    ).images[0]
    return image

# Gradio Interface
demo = gr.Interface(
    fn=generate,
    inputs=[
        gr.Textbox(
            label="Prompt",
            value="product photo of Indian Saree, black background, no person, studio shot"
        ),
        gr.Textbox(
            label="Negative Prompt",
            value="human face, person, human body, skin texture, portrait"
        ),
        gr.Slider(10, 50, value=30, step=1, label="Steps"),
        gr.Slider(1, 15, value=7.5, step=0.5, label="Guidance Scale"),
    ],
    outputs=gr.Image(label="Generated Image"),
    title="🇮🇳 Indian Clothing Generator",
    description="Generate product photos of Indian clothing (Saree, Kurta, Shirt, Jacket, T-shirt)",
    examples=[
        ["product photo of Indian Saree, black background, no person", "human face, person", 30, 7.5],
        ["product photo of Indian Kurta, black background, no person", "human face, person", 30, 7.5],
        ["product photo of Indian Jacket, black background, no person", "human face, person", 30, 7.5],
    ]
)

demo.launch()