#!/usr/bin/env python3
"""
Image-to-Image Variation App mit Stable Diffusion und Gradio
Ermöglicht es, Variationen von bestehenden Bildern zu generieren
"""

import gradio as gr
import torch
from diffusers import StableDiffusionImg2ImgPipeline
from PIL import Image
import numpy as np

# Gerät konfigurieren (GPU falls verfügbar)
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")

# Modell laden
print("Loading Stable Diffusion model...")
pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
    safety_checker=None
)
pipe = pipe.to(device)

# Optional: Enable memory optimizations for slower devices
if device == "cpu":
    pipe.enable_attention_slicing()

def generate_image_variation(input_image, prompt, negative_prompt, strength, guidance_scale, num_steps):
    """
    Generiert eine Variation des Eingangsbildes basierend auf dem Prompt
    
    Args:
        input_image: PIL Image oder numpy array
        prompt: Text-Beschreibung der gewünschten Variation
        negative_prompt: Was nicht im Bild sein soll
        strength: Wie stark die Veränderung sein soll (0.0-1.0)
        guidance_scale: Wie stark der Prompt befolgt werden soll
        num_steps: Anzahl der Diffusion-Schritte
    
    Returns:
        Generiertes Variations-Bild
    """
    
    # Input validieren
    if input_image is None:
        raise gr.Error("Bitte laden Sie zuerst ein Bild hoch")
    
    if not prompt.strip():
        raise gr.Error("Bitte geben Sie einen Text-Prompt ein")
    
    # Bild vorbereiten
    if isinstance(input_image, np.ndarray):
        input_image = Image.fromarray(input_image)
    
    # Auf einheitliche Größe skalieren (für bessere Performance)
    input_image = input_image.convert("RGB")
    input_image = input_image.resize((512, 512), Image.Resampling.LANCZOS)
    
    print(f"Generating variation with prompt: {prompt}")
    print(f"Strength: {strength}, Guidance: {guidance_scale}, Steps: {num_steps}")
    
    # Generiere Variation
    with torch.no_grad():
        result = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt if negative_prompt.strip() else None,
            image=input_image,
            strength=strength,
            guidance_scale=guidance_scale,
            num_inference_steps=num_steps,
            height=512,
            width=512,
            generator=torch.Generator(device).manual_seed(42)  # Für Reproduzierbarkeit
        )
    
    return result.images[0]

# Gradio Interface erstellen
with gr.Blocks(title="Image-to-Image Variation") as demo:
    gr.Markdown("# 🎨 Image-to-Image Variation mit Stable Diffusion")
    gr.Markdown("Generieren Sie kreative Variationen von Ihren Bildern durch Text-Prompts")
    
    with gr.Row():
        with gr.Column():
            gr.Markdown("### Eingabe")
            
            # Input Bild
            input_image = gr.Image(
                label="Eingabe-Bild",
                type="pil",
                sources=["upload", "webcam"],
                height=400
            )
            
            # Text Prompts
            prompt = gr.Textbox(
                label="Prompt",
                placeholder="z.B. 'Van Gogh Stil', 'Ölgemälde', 'Realistic photograph'",
                lines=3
            )
            
            negative_prompt = gr.Textbox(
                label="Negativer Prompt (optional)",
                placeholder="z.B. 'blurry, low quality, distorted'",
                lines=2
            )
            
            with gr.Row():
                strength = gr.Slider(
                    label="Strength (Veränderungsstärke)",
                    minimum=0.0,
                    maximum=1.0,
                    value=0.7,
                    step=0.05,
                    info="Höher = mehr Veränderung vom Original"
                )
                
                guidance_scale = gr.Slider(
                    label="Guidance Scale",
                    minimum=1.0,
                    maximum=20.0,
                    value=7.5,
                    step=0.5,
                    info="Wie stark der Prompt befolgt wird"
                )
            
            steps = gr.Slider(
                label="Inference Steps",
                minimum=20,
                maximum=100,
                value=50,
                step=5,
                info="Mehr = bessere Qualität, aber länger"
            )
            
            # Generate Button
            generate_btn = gr.Button("🎨 Variation generieren", variant="primary", size="lg")
        
        with gr.Column():
            gr.Markdown("### Ausgabe")
            output_image = gr.Image(
                label="Generiertes Bild",
                type="pil",
                height=400
            )
    
    # Examples
    gr.Markdown("### Beispiele zum Ausprobieren:")
    
    examples_data = [
        ["example_dog.png", "Van Gogh painting style", "", 0.7, 7.5, 50],
        ["example_landscape.png", "Oil painting, impressionist", "", 0.75, 7.5, 50],
        ["example_portrait.png", "Anime art style", "low quality, blurry", 0.6, 7.5, 50],
    ]
    
    gr.Examples(
        examples=[
            [None, "A beautiful landscape in watercolor style", "", 0.7, 7.5, 50],
            [None, "Oil painting, renaissance style", "", 0.75, 7.5, 50],
            [None, "Modern digital art, cyberpunk aesthetic", "blurry, low quality", 0.6, 7.5, 50],
        ],
        inputs=[input_image, prompt, negative_prompt, strength, guidance_scale, steps],
        outputs=output_image,
        fn=generate_image_variation,
        cache_examples=False,
        run_on_click=False
    )
    
    # Infos
    with gr.Row():
        with gr.Column():
            gr.Markdown("""
            #### 💡 Tipps:
            - **Strength**: 0.3-0.5 für subtile Änderungen, 0.7-0.9 für drastische Veränderungen
            - **Guidance Scale**: 7.5-10 für gute Balance zwischen Qualität und Prompt-Treue
            - **Steps**: 30-50 normalerweise ausreichend, mehr für höhere Details
            """)
        
        with gr.Column():
            gr.Markdown(f"""
            #### ℹ️ Info:
            - **Device**: {device.upper()}
            - **Model**: Stable Diffusion v1.5
            - **Input size**: 512x512
            - **Format**: PNG, JPG, WebP
            """)
    
    # Click handler
    generate_btn.click(
        fn=generate_image_variation,
        inputs=[input_image, prompt, negative_prompt, strength, guidance_scale, steps],
        outputs=output_image
    )

if __name__ == "__main__":
    # App starten
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        show_error=True,
        debug=True
    )