import gradio as gr
from diffusers import AutoPipelineForText2Image
import torch
import gc

# --------------------------------------------------------
# 1. Loading the Turbo Model for CPU Inference
# --------------------------------------------------------
print("🔄 Loading the Text-to-Image AI Model (SD Turbo - Fast Inference)...")

# We use the AutoPipelineForText2Image from diffusers.
# 'stabilityai/sd-turbo' allows acceptable quality in just 1 to 4 steps!
# This makes it viable for CPU-only execution in HuggingFace free spaces.
try:
    pipe = AutoPipelineForText2Image.from_pretrained(
        "stabilityai/sd-turbo",
        torch_dtype=torch.float32, 
        variant="fp16" # use fp16 weights where possible to save memory
    )
    pipe = pipe.to("cpu")
    
    # Optional performance tweaks for CPU:
    pipe.set_progress_bar_config(disable=True)
    
    print("✅ Model loaded successfully!")
    print(f"Pipeline components: {list(pipe.components.keys())}")
except Exception as e:
    print(f"❌ Error loading model: {e}")
    # Fallback in case fp16 variant fails to download on CPU
    pipe = AutoPipelineForText2Image.from_pretrained(
        "stabilityai/sd-turbo",
        torch_dtype=torch.float32
    )
    pipe = pipe.to("cpu")
    print("✅ Model loaded via fallback!")

# --------------------------------------------------------
# 2. Generation Logic
# --------------------------------------------------------
def generate_image(prompt, num_steps, guidance_scale, seed):
    if not prompt or not prompt.strip():
        raise gr.Error("⚠️ Please enter a text prompt.")
        
    print(f"Generating: '{prompt}' (Steps: {num_steps}, Seed: {seed})")
    
    # Manage seed reproducibility
    generator = torch.Generator("cpu").manual_seed(int(seed))
    
    try:
        # Generate the image 
        # SD-Turbo performs best around 1-4 steps
        result = pipe(
            prompt=prompt,
            num_inference_steps=int(num_steps),
            guidance_scale=float(guidance_scale), # Usually 0.0 for SD-Turbo
            generator=generator
        )
        
        # Free up memory
        gc.collect()
        
        return result.images[0]
        
    except Exception as e:
        import traceback
        traceback.print_exc()
        raise gr.Error(f"❌ Error during generation: {str(e)}")

# --------------------------------------------------------
# 3. Custom UI Styling
# --------------------------------------------------------
custom_css = """
@import url('https://fonts.googleapis.com/css2?family=Outfit:wght@300;400;600;800&display=swap');

* {
    font-family: 'Outfit', sans-serif !important;
}

.gradio-container {
    max-width: 1000px !important;
    margin: auto !important;
    background: radial-gradient(circle at 50% 0%, #1e293b 0%, #0f172a 100%) !important;
    min-height: 100vh;
}

.header-container {
    text-align: center;
    padding: 30px;
    margin-bottom: 20px;
    border-radius: 20px;
    background: rgba(255, 255, 255, 0.03);
    border: 1px solid rgba(255, 255, 255, 0.1);
    box-shadow: 0 10px 30px rgba(0,0,0,0.5);
    backdrop-filter: blur(10px);
}

.title-text {
    font-size: 3rem !important;
    font-weight: 800 !important;
    background: linear-gradient(135deg, #00f2fe 0%, #4facfe 100%) !important;
    -webkit-background-clip: text !important;
    -webkit-text-fill-color: transparent !important;
    margin-bottom: 15px !important;
    letter-spacing: -1px;
}

.subtitle-text {
    color: #94a3b8 !important;
    font-size: 1.2rem !important;
    font-weight: 300 !important;
}

.generate-btn {
    background: linear-gradient(135deg, #00f2fe 0%, #4facfe 100%) !important;
    border: none !important;
    box-shadow: 0 4px 15px rgba(79, 172, 254, 0.4) !important;
    color: white !important;
    font-weight: 800 !important;
    font-size: 1.2rem !important;
    border-radius: 12px !important;
    transition: all 0.3s ease !important;
    padding: 15px !important;
}

.generate-btn:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 8px 25px rgba(79, 172, 254, 0.6) !important;
}

/* Make image display beautiful */
.image-output img {
    border-radius: 12px !important;
    box-shadow: 0 10px 25px rgba(0,0,0,0.5) !important;
}
"""

# --------------------------------------------------------
# 4. Gradio Application Construction
# --------------------------------------------------------
with gr.Blocks(css=custom_css, title="🎨 Fast Text-to-Image AI", theme=gr.themes.Monochrome()) as demo:
    
    # Header Section
    gr.HTML("""
    <div class="header-container">
        <h1 class="title-text">⚡ Fast Text-to-Image AI</h1>
        <p class="subtitle-text">Powered by SD-Turbo. Generates beautiful images on CPU in seconds.</p>
    </div>
    """)
    
    with gr.Row():
        # Left Column - Controls
        with gr.Column(scale=1):
            prompt = gr.Textbox(
                label="🔮 Your Prompt",
                placeholder="A futuristic city at sunset, highly detailed, cyberpunk style...",
                lines=3
            )
            
            with gr.Accordion("⚙️ Advanced Settings", open=False):
                num_steps = gr.Slider(
                    label="Steps (Quality vs Speed)",
                    minimum=1, 
                    maximum=10, 
                    value=2, 
                    step=1,
                    info="SD-Turbo is designed for 1-4 steps!"
                )
                guidance_scale = gr.Slider(
                    label="Guidance Scale",
                    minimum=0.0, 
                    maximum=5.0, 
                    value=0.0, 
                    step=0.1,
                    info="Must be 0.0 for SD-Turbo!"
                )
                seed = gr.Slider(
                    label="Random Seed",
                    minimum=1, 
                    maximum=999999, 
                    value=1337, 
                    step=1,
                    info="Change to get different images"
                )
                
            generate_btn = gr.Button("🚀 Generate Image", elem_classes=["generate-btn"])
            
            gr.Examples(
                examples=[
                    ["A cute corgi dog in a spacesuit on Mars", 2, 0.0, 42],
                    ["A hyper-realistic photograph of a juicy hamburger with melted cheese", 3, 0.0, 100],
                    ["Cinematic shot of an ancient mechanical dragon sleeping in a cave", 4, 0.0, 999]
                ],
                inputs=[prompt, num_steps, guidance_scale, seed],
                label="💡 Try examples"
            )
            
        # Right Column - Output
        with gr.Column(scale=1):
            output_image = gr.Image(
                label="✨ Generated Result", 
                type="pil", 
                elem_classes=["image-output"]
            )
            
    # Connect UI to logic
    generate_btn.click(
        fn=generate_image,
        inputs=[prompt, num_steps, guidance_scale, seed],
        outputs=output_image
    )

# Launch app
demo.launch()