from diffusers import StableDiffusionPipeline
import gradio as gr
import torch


# Load model
#model_id = "stabilityai/stable-diffusion-2-1"
model_id = "runwayml/stable-diffusion-v1-5"
device = "cuda" if torch.cuda.is_available() else "cpu"

pipe = StableDiffusionPipeline.from_pretrained(
    model_id,
    torch_dtype=torch.float16 if device == "cuda" else torch.float32
).to(device)


# Function to generate mockup image from text prompt
def generate_mockup(prompt, num_inference_steps=50, guidance_scale=7.5):
    if device == "cuda":
        image = pipe(
            prompt,
            num_inference_steps=num_inference_steps,
            guidance_scale=guidance_scale
        ).images[0]
    else:
        with torch.no_grad():
            image = pipe(
                prompt,
                num_inference_steps=num_inference_steps,
                guidance_scale=guidance_scale
            ).images[0]

    # Simple info text under the image
    info = f"Image Size: {image.size[0]}×{image.size[1]} — Model: Stable Diffusion-v1-5"

    return image, info

# Custom Theme
theme = gr.themes.Soft(
    primary_hue="pink",
    secondary_hue="gray",
    neutral_hue="slate",
    font=["Inter", "sans-serif"],
).set(
    button_primary_background_fill="linear-gradient(90deg, #6366F1, #3B82F6)",
    button_primary_text_color="white",
    button_primary_background_fill_hover="linear-gradient(90deg, #4F46E5, #2563EB)",
    button_border_width="0px",
)

# Interface
with gr.Blocks(theme=theme, css="""
#title {text-align:center; font-size:1.8em; font-weight:700; margin-bottom:0.4em;}
#subtitle {text-align:center; color:#6B7280; font-size:1em; margin-bottom:1.2em;}
.footer {text-align:center; font-size:0.85em; color:#9CA3AF; margin-top:1.5em;}
""") as demo:

    gr.Markdown("<div id='title'>AI Product Mockup Generator</div>")
    gr.Markdown("<div id='subtitle'>Turn your text ideas into professional product mockups instantly!</div>")

    with gr.Row():
        with gr.Column(scale=1):
            prompt = gr.Textbox(
                label="Enter your product description",
                placeholder="e.g., A modern smartwatch on a white background",
                lines=2
            )

            # Advanced settings hidden under an accordion
            with gr.Accordion("Advanced Settings", open=False):
                steps = gr.Slider(1, 50, value=10, step=5, label="Inference Steps")
                guidance = gr.Slider(1.0, 15.0, value=7.5, step=0.1, label="Guidance Scale")

            with gr.Row():
                generate_btn = gr.Button("Generate Mockup", variant="primary")
                clear_btn = gr.Button("Clear")

        with gr.Column(scale=1):
            output_image = gr.Image(type="pil", label="Generated Mockup", show_label=False)
            image_info = gr.Markdown("")

    # Button functions
    generate_btn.click(generate_mockup, [prompt, steps, guidance], [output_image, image_info])

    # Clear button resets everything
    clear_btn.click(lambda: ("", None, ""), None, [prompt, output_image, image_info])

    #gr.Markdown("<div class='footer'>Created with Stable Diffusion 2.1 and Gradio</div>")

# Launch App
demo.launch()