import spaces
import torch
from diffusers import StableDiffusionPipeline
import gradio as gr
import random
from PIL import Image
import os

# Initialisiere die Pipeline als None - wird erst bei Bedarf geladen
pipe = None

def load_model():
    global pipe
    if pipe is None:
        print("Loading SeaCreatures model...")
        try:
            pipe = StableDiffusionPipeline.from_pretrained(
                "arnomatic/seacreatures",
                torch_dtype=torch.float16,
                safety_checker=None,
                requires_safety_checker=False,
                use_safetensors=True
            )
            print("Model loaded successfully!")
        except Exception as e:
            print(f"Error loading model: {e}")
            # Fallback zu Standard Stable Diffusion falls dein Modell nicht lädt
            pipe = StableDiffusionPipeline.from_pretrained(
                "runwayml/stable-diffusion-v1-5",
                torch_dtype=torch.float16,
                safety_checker=None,
                requires_safety_checker=False
            )
            print("Loaded fallback model instead")
    return pipe

# Beispiel-Prompts für Inspiration
example_prompts = [
    "a majestic jellyfish with glowing tentacles",
    "an ethereal sea dragon with translucent fins",
    "a mysterious deep sea creature with bioluminescent spots", 
    "an elegant manta ray with glowing patterns",
    "a fantastical seahorse with luminous spines",
    "a graceful octopus with shimmering skin",
    "a mystical anglerfish with bright lure",
    "an otherworldly nautilus with glowing shell patterns"
]

@spaces.GPU(duration=120)  # 2 Minuten sollten reichen
def generate_seacreature(prompt, negative_prompt, num_inference_steps, guidance_scale, width, height, seed):
    # Lade das Modell falls noch nicht geschehen
    current_pipe = load_model()
    
    # Setze das Modell auf GPU
    current_pipe.to("cuda")
    
    # Füge den Concept-Trigger hinzu - das ist wichtig für dein DreamBooth Modell!
    full_prompt = f"underwater bioluminescence creature, {prompt}"
    
    # Seed handling
    if seed == -1:
        seed = random.randint(0, 2**32 - 1)
    generator = torch.Generator(device="cuda").manual_seed(seed)
    
    try:
        # Generiere das Bild
        with torch.no_grad():
            result = current_pipe(
                prompt=full_prompt,
                negative_prompt=negative_prompt,
                num_inference_steps=int(num_inference_steps),
                guidance_scale=guidance_scale,
                width=int(width),
                height=int(height),
                generator=generator
            )
            
        image = result.images[0]
        return image, seed
        
    except Exception as e:
        print(f"Error during generation: {e}")
        # Fallback: Erstelle ein einfaches Bild mit Fehlermeldung
        error_img = Image.new('RGB', (int(width), int(height)), color='navy')
        return error_img, seed
    finally:
        # GPU wieder freigeben
        current_pipe.to("cpu")
        torch.cuda.empty_cache()

def get_random_prompt():
    return random.choice(example_prompts)

# CSS für schönere Optik
css = """
.gradio-container {
    font-family: 'Helvetica Neue', Arial, sans-serif;
}
.title {
    text-align: center;
    color: #0077be;
    margin-bottom: 20px;
}
.description {
    text-align: center;
    margin-bottom: 20px;
    color: #666;
}
"""

# Gradio Interface
with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
    gr.HTML("<h1 class='title'>🌊 SeaCreatures - Bioluminescent Underwater Beings ✨</h1>")
    gr.HTML("<p class='description'>Generate mystical underwater creatures with bioluminescent features using a fine-tuned Stable Diffusion model.</p>")
    
    with gr.Row():
        with gr.Column(scale=1):
            prompt = gr.Textbox(
                label="Prompt",
                placeholder="Describe your underwater creature...",
                lines=3,
                value="a majestic jellyfish with glowing tentacles"
            )
            
            negative_prompt = gr.Textbox(
                label="Negative Prompt (Optional)",
                placeholder="What you don't want to see...",
                lines=2,
                value="blurry, low quality, distorted, ugly, bad anatomy"
            )
            
            with gr.Row():
                random_prompt_btn = gr.Button("🎲 Random Prompt", size="sm")
                generate_btn = gr.Button("🎨 Generate SeaCreature", variant="primary", size="lg")
            
            with gr.Accordion("Advanced Settings", open=False):
                num_inference_steps = gr.Slider(
                    label="Inference Steps",
                    minimum=10,
                    maximum=50,
                    step=5,
                    value=25,
                    info="More steps = better quality but slower"
                )
                
                guidance_scale = gr.Slider(
                    label="Guidance Scale",
                    minimum=1.0,
                    maximum=20.0,
                    step=0.5,
                    value=7.5,
                    info="Higher values follow prompt more closely"
                )
                
                with gr.Row():
                    width = gr.Slider(
                        label="Width",
                        minimum=256,
                        maximum=768,
                        step=64,
                        value=512
                    )
                    height = gr.Slider(
                        label="Height", 
                        minimum=256,
                        maximum=768,
                        step=64,
                        value=512
                    )
                
                seed = gr.Number(
                    label="Seed (-1 for random)",
                    value=-1,
                    precision=0
                )
        
        with gr.Column(scale=1):
            output_image = gr.Image(
                label="Generated SeaCreature",
                type="pil",
                height=512
            )
            
            used_seed = gr.Number(
                label="Used Seed",
                interactive=False
            )
    
    # Example Gallery
    gr.HTML("<h3 style='text-align: center; margin-top: 30px;'>💡 Example Prompts:</h3>")
    
    examples = gr.Examples(
        examples=[
            ["a majestic jellyfish with glowing tentacles", "blurry, low quality", 25, 7.5, 512, 512, -1],
            ["an ethereal sea dragon with translucent fins", "blurry, low quality", 25, 7.5, 512, 512, -1],
            ["a mysterious deep sea creature with bioluminescent spots", "blurry, low quality", 25, 7.5, 512, 512, -1],
            ["an elegant manta ray with glowing patterns", "blurry, low quality", 25, 7.5, 512, 512, -1],
        ],
        inputs=[prompt, negative_prompt, num_inference_steps, guidance_scale, width, height, seed],
        outputs=[output_image, used_seed],
        fn=generate_seacreature,
        cache_examples=False
    )
    
    # Event handlers
    random_prompt_btn.click(
        fn=get_random_prompt,
        outputs=prompt
    )
    
    generate_btn.click(
        fn=generate_seacreature,
        inputs=[prompt, negative_prompt, num_inference_steps, guidance_scale, width, height, seed],
        outputs=[output_image, used_seed]
    )

# Info Footer
    gr.HTML("""
    <div style='text-align: center; margin-top: 30px; padding: 20px; background-color: #f0f8ff; border-radius: 10px;'>
        <p><strong>💡 Tips:</strong></p>
        <p>• The model automatically adds "underwater bioluminescence creature" to your prompt</p>
        <p>• Try describing colors, shapes, and lighting effects for best results</p>
        <p>• Use negative prompts to avoid unwanted elements</p>
        <p>• Model created by <a href="https://huggingface.co/arnomatic" target="_blank">@arnomatic</a></p>
    </div>
    """)

if __name__ == "__main__":
    demo.launch()