import gradio as gr
import torch
from diffusers import StableDiffusionPipeline
from peft import PeftModel
import spaces

MODEL_REPO = "DAVEinside/Vintage_art_LORA"

@spaces.GPU
def load_pipeline():
    """Load the pipeline with LoRA weights"""
    try:
        pipe = StableDiffusionPipeline.from_pretrained(
            "runwayml/stable-diffusion-v1-5",
            torch_dtype=torch.float16,
            safety_checker=None,
            requires_safety_checker=False
        )
        
        # Load LoRA weights
        pipe.unet = PeftModel.from_pretrained(pipe.unet, MODEL_REPO)
        pipe = pipe.to("cuda")
        
        return pipe
    except Exception as e:
        print(f"Error loading pipeline: {e}")
        return None

# Load pipeline globally
print("🔄 Loading Vintage Travel Poster LoRA...")
pipe = load_pipeline()

@spaces.GPU
def generate_poster(destination, style_elements, custom_prompt, steps, guidance):
    """Generate vintage travel poster"""
    
    if pipe is None:
        return None, "❌ Model failed to load"
    
    try:
        # Build prompt
        if custom_prompt.strip():
            prompt = custom_prompt
        else:
            prompt = f"vintage travel poster of {destination}"
            if style_elements:
                prompt += f", {style_elements}"
            prompt += ", art deco style, tourism poster"
        
        # Negative prompt for better quality
        negative_prompt = "blurry, low quality, distorted, watermark, signature, text, words, letters"
        
        print(f"🎨 Generating: {prompt}")
        
        # Generate image
        image = pipe(
            prompt,
            negative_prompt=negative_prompt,
            num_inference_steps=int(steps),
            guidance_scale=float(guidance),
            width=512,
            height=512
        ).images[0]
        
        return image, f"✅ Generated: {prompt}"
        
    except Exception as e:
        return None, f"❌ Generation failed: {str(e)}"

# Create Gradio interface
with gr.Blocks(title="🎨 Vintage Travel Posters", theme=gr.themes.Soft()) as app:
    gr.HTML("""
    <div style="text-align: center; margin-bottom: 20px;">
        <h1>🎨 Vintage Travel Posters Generator</h1>
        <p style="font-size: 18px;">Transform any destination into beautiful vintage travel poster art!</p>
        <p>Create stunning retro-style posters in the classic 1930s-1950s tourism style.</p>
    </div>
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### ✏️ Design Your Poster")
            
            destination = gr.Textbox(
                label="🗺️ Destination",
                placeholder="e.g., Paris, Tokyo, New York, Swiss Alps...",
                value="Paris"
            )
            
            style_elements = gr.Textbox(
                label="🎭 Style Elements",
                placeholder="e.g., Eiffel Tower, sunset, mountains, beaches...",
                value="Eiffel Tower, romantic, golden hour"
            )
            
            custom_prompt = gr.Textbox(
                label="✨ Custom Prompt (Optional)",
                placeholder="Write your own complete prompt here...",
                lines=3
            )
            
            with gr.Row():
                steps = gr.Slider(
                    minimum=10, 
                    maximum=30, 
                    value=20, 
                    step=1, 
                    label="🔧 Inference Steps"
                )
                guidance = gr.Slider(
                    minimum=5.0, 
                    maximum=15.0, 
                    value=7.5, 
                    step=0.5, 
                    label="🎯 Guidance Scale"
                )
            
            generate_btn = gr.Button("🎨 Generate Vintage Poster", variant="primary", size="lg")
            
            gr.Markdown("""
            ### 💡 Tips for Best Results:
            - **Destinations**: Cities, countries, landmarks work great
            - **Elements**: Include iconic landmarks, colors, activities  
            - **Styles**: Try "art deco", "1940s style", "tourism poster"
            - **Trigger words**: "vintage travel poster", "retro illustration"
            """)
            
        with gr.Column(scale=1):
            gr.Markdown("### 🖼️ Your Vintage Poster")
            
            output_image = gr.Image(
                label="Generated Poster",
                height=512,
                width=512
            )
            
            status_text = gr.Textbox(
                label="Status",
                interactive=False,
                lines=2
            )
    
    # Example gallery
    gr.Markdown("### 🎯 Example Prompts")
    
    examples = [
        ["Paris", "Eiffel Tower, romantic, golden hour", "", 20, 7.5],
        ["New York", "skyline, skyscrapers, art deco", "", 20, 7.5],
        ["Tokyo", "Mount Fuji, cherry blossoms, traditional", "", 20, 7.5],
        ["California", "palm trees, beaches, sunshine, surfing", "", 20, 7.5],
        ["Swiss Alps", "mountains, skiing, winter sports, snow", "", 20, 7.5],
        ["London", "Big Ben, red buses, fog, British", "", 20, 7.5],
        ["", "", "vintage travel poster of Rome, Colosseum, warm colors, 1940s style", 20, 7.5],
        ["", "", "vintage travel poster of Hawaii, tropical beaches, surfing, paradise", 20, 7.5]
    ]
    
    gr.Examples(
        examples=examples,
        inputs=[destination, style_elements, custom_prompt, steps, guidance],
        outputs=[output_image, status_text],
        fn=generate_poster,
        cache_examples=False,
        label="Click any example to try it:"
    )
    
    # Event handlers
    generate_btn.click(
        fn=generate_poster,
        inputs=[destination, style_elements, custom_prompt, steps, guidance],
        outputs=[output_image, status_text]
    )
    
    gr.HTML("""
    <div style="text-align: center; margin-top: 30px; padding: 20px; background-color: #f8f9fa; border-radius: 10px;">
        <h3>🎨 About This Model</h3>
        <p>This demo uses a custom-trained LoRA model that specializes in vintage travel poster artwork.</p>
        <p>Perfect for designers, marketers, and art enthusiasts who love the classic travel poster aesthetic!</p>
        <hr style="margin: 20px 0;">
        <p><strong>🏷️ Best Trigger Words:</strong> "vintage travel poster", "art deco style", "tourism poster", "1940s style"</p>
        <p><strong>🎯 Model:</strong> Stable Diffusion v1.5 + Custom LoRA</p>
        <p><strong>🎨 Style:</strong> 1930s-1950s Vintage Tourism Posters</p>
    </div>
    """)

# Launch the app
if __name__ == "__main__":
    app.launch()