import gradio as gr
import torch
from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
from PIL import Image
import random
import os

# Model selection based on art style
MODEL_MAP = {
    "default": "stabilityai/stable-diffusion-xl-base-1.0",
    "anime": "cagliostrolab/animagine-xl-3.1",
    "pixel": "wavymulder/pixel-art-diffusion",
    "cinematic": "stabilityai/stable-diffusion-3.5-medium",
}

class PerchanceStyleGenerator:
    def __init__(self):
        self.current_model = None
        self.pipe = None
        
    def load_model(self, model_key="default"):
        if self.current_model != model_key:
            model_id = MODEL_MAP.get(model_key, MODEL_MAP["default"])
            print(f"Loading {model_id}...")
            self.pipe = StableDiffusionXLPipeline.from_pretrained(
                model_id,
                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
                use_safetensors=True
            )
            self.pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(self.pipe.scheduler.config)
            
            if torch.cuda.is_available():
                self.pipe = self.pipe.to("cuda")
            self.current_model = model_key
    
    def generate(self, prompt, negative, style, shape, num_images):
        # Map Perchance style to model
        style_lower = style.lower()
        if any(anime in style_lower for anime in ["anime", "manga", "waifu"]):
            model_key = "anime"
        elif "pixel" in style_lower:
            model_key = "pixel"
        elif any(cine in style_lower for cine in ["cinematic", "photo"]):
            model_key = "cinematic"
        else:
            model_key = "default"
        
        self.load_model(model_key)
        
        # Calculate dimensions based on shape
        if shape == "Square":
            width, height = 1024, 1024
        elif shape == "Portrait":
            width, height = 832, 1216
        else:  # Landscape
            width, height = 1216, 832
        
        # Enhance prompt with style
        style_prompt = f"{prompt}, {style.lower()} style, high quality, detailed"
        
        # Generate images
        images = []
        for i in range(num_images):
            generator = torch.Generator().manual_seed(random.randint(1, 999999))
            
            result = self.pipe(
                prompt=style_prompt,
                negative_prompt=negative,
                width=width,
                height=height,
                num_inference_steps=25,
                guidance_scale=7.5,
                generator=generator
            )
            images.append(result.images[0])
        
        if len(images) == 1:
            return images[0]
        else:
            return self.make_grid(images, cols=min(4, num_images))
    
    def make_grid(self, images, cols=4):
        rows = (len(images) + cols - 1) // cols
        w, h = images[0].size
        grid = Image.new('RGB', (cols * w, rows * h))
        for i, img in enumerate(images):
            grid.paste(img, (i % cols * w, i // cols * h))
        return grid

generator = PerchanceStyleGenerator()

ART_STYLES = [
    "Painted", "Anime", "Casual Photo", "Cinematic", "Digital Painting", 
    "Concept Art", "No style", "3D", "Disney Character", "2D Disney Character", 
    "Disney Sketch", "Concept Sketch", "Painterly", "Oil Painting", 
    "Oil Painting - Realism", "Oil Painting - Old", "Oil Painting - 70s Pulp", 
    "Professional Photo", "Anime Drawn", "Anime Anime", "Anime Screencap", 
    "Cute Anime", "Soft Anime", "Fantasy Painting", "Fantasy Landscape", 
    "Fantasy Portrait", "Studio Ghibli", "50s Enamel Sign", "Vintage Comic", 
    "Franco-Belgian Comic", "Tintin Comic", "Medieval", "Pixel Art", 
    "Furry - Oil", "Furry - Cinematic", "Furry - Painted", "Furry - Drawn", 
    "Cute Figurine", "3D Emoji", "Illustration", "Cute Illustration", 
    "Flat Illustration", "Watercolor", "1990s Photo", "1980s Photo", 
    "1970s Photo", "1960s Photo", "1950s Photo", "1940s Photo", "1930s Photo", 
    "1920s Photo", "Vintage Pulp Art", "50s Infomercial", "Anime 3D", 
    "Pokemon Painted", "2D Pokemon", "Vintage Anime", "Neon Vintage Anime", 
    "Manga", "Fantasy World Map", "Fantasy City Map", "Old World Map", 
    "3D Isometric", "Icon", "Flat Style Icon", "Flat Style Logo", 
    "Game Art Icon", "Digital Painting Icon", "Concept Art Icon", 
    "Cute 3D Icon", "Cute 3D Icon Set", "Crayon Drawing", "Pencil", 
    "Tattoo Design", "Waifu", "YuGiOh Art", "Traditional Japanese", 
    "Nihonga Painting", "Claymation", "Cartoon", "Cursed Photo", "MTG Card"
]

with gr.Blocks(theme=gr.themes.Soft(), title="ImageMagic") as demo:
    gr.Markdown("# 🌌 ImageMagic")
    gr.Markdown("Like Perchance, but self-hosted on Hugging Face!")
    
    with gr.Row():
        with gr.Column(scale=1):
            prompt = gr.Textbox(
                label="💭 Description",
                placeholder="Describe what you want to see...",
                lines=3
            )
            
            negative = gr.Textbox(
                label="🚫 Anti-Description (optional)",
                placeholder="What you DON'T want...",
                lines=2
            )
            
            style = gr.Dropdown(
                label="🎨 Art Style",
                choices=ART_STYLES,
                value="Cinematic"
            )
            
            with gr.Row():
                shape = gr.Radio(
                    label="🖼️ Shape",
                    choices=["Square", "Portrait", "Landscape"],
                    value="Square"
                )
                
                num_images = gr.Dropdown(
                    label="🔢 How many?",
                    choices=[1, 2, 4, 6, 8],
                    value=1
                )
            
            generate_btn = gr.Button("✨ Generate", variant="primary", size="lg")
        
        with gr.Column(scale=1):
            output = gr.Image(label="Result", height=500)
    
    generate_btn.click(
        fn=generator.generate,
        inputs=[prompt, negative, style, shape, num_images],
        outputs=output
    )

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