import gradio as gr
import torch
from diffusers import DiffusionPipeline

# Load the new model
# Use bfloat16 for speed and CUDA if available
device = "cuda" if torch.cuda.is_available() else "cpu"
pipe = DiffusionPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev", 
    dtype=torch.bfloat16,
    device_map=device
)

# Function to generate outfit image + suggestion
def generate_outfit(weather, activity, style):
    try:
        prompt = (
            f"A realistic full-body outfit for {weather} weather, "
            f"{activity} activity, {style} style, front view, fashion style, full body"
        )
        result = pipe(prompt)
        image = result.images[0]

        suggestion_text = (
            f"💡 Suggested outfit for {weather} weather, {activity} activity, {style} style:\n"
            "- Top: light breathable shirt or jacket depending on weather\n"
            "- Bottom: comfortable pants or shorts\n"
            "- Shoes: suitable for activity\n"
            "- Accessories: hat, sunglasses, or scarf depending on weather"
        )
        return image, suggestion_text
    except Exception as e:
        return None, f"❌ Error generating outfit: {str(e)}"

# Gradio UI
custom_css = """
body {background: linear-gradient(135deg, #232526, #414345); color: white; font-family: 'Poppins', sans-serif;}
.gradio-container {max-width: 900px !important; margin: auto;}
h1,h2,h3 {text-align:center; color:#fff;}
.gr-button {background-color:#00bfa6 !important; color:white !important; font-weight:bold; border-radius:12px !important;}
.gr-button:hover {background-color:#007f70 !important;}
.gr-image {border-radius:12px; box-shadow:0 0 20px rgba(0,255,255,0.1);}
"""

with gr.Blocks(css=custom_css) as app:
    gr.Markdown("<h1 style='text-align:center;'>👗 ClothCast - AI Outfit Forecaster (Fast)</h1>")
    gr.Markdown("<p style='text-align:center;'>Select your weather, activity, and style to generate an AI outfit quickly!</p>")

    with gr.Row():
        with gr.Column(scale=1):
            weather = gr.Dropdown(
                ["Hot", "Cold", "Rainy", "Snowy", "Mild", "Humid"],
                label="🌤 Weather", value="Hot"
            )
            activity = gr.Dropdown(
                ["Casual", "Work", "Party", "Sporty"],
                label="🏃 Activity", value="Casual"
            )
            style = gr.Dropdown(
                ["Casual", "Sporty", "Formal"],
                label="🎨 Style", value="Casual"
            )
            generate_btn = gr.Button("Generate Outfit 👗", variant="primary")

        with gr.Column(scale=1):
            outfit_image = gr.Image(label="🖼 Generated Outfit")
            outfit_text = gr.Textbox(label="💡 Outfit Suggestion", lines=8)

    generate_btn.click(
        generate_outfit, 
        inputs=[weather, activity, style], 
        outputs=[outfit_image, outfit_text]
    )

app.launch(share=True)