import gradio as gr
import torch
from diffusers import StableDiffusionPipeline
from huggingface_hub import hf_hub_download
import numpy as np
from PIL import Image
import os

# Suppress symlink warnings
os.environ['HF_HUB_DISABLE_SYMLINKS_WARNING'] = "1"

# Define styles
styles = {
    "glitch": {
        "concept_url": "sd-concepts-library/001glitch-core",
        "seed": 42,
        "token": "<glitch-core>"
    },
    "roth": {
        "concept_url": "sd-concepts-library/2814-roth",
        "seed": 123,
        "token": "<2814-roth>"
    },
    "night": {
        "concept_url": "sd-concepts-library/4tnght",
        "seed": 456,
        "token": "<4tnght>"
    },
    "anime80s": {
        "concept_url": "sd-concepts-library/80s-anime-ai",
        "seed": 789,
        "token": "<80s-anime>"
    },
    "animeai": {
        "concept_url": "sd-concepts-library/80s-anime-ai-being",
        "seed": 1024,
        "token": "<80s-anime-being>"
    }
}

def load_pipeline():
    """Load and prepare the pipeline with all style embeddings"""
    pipe = StableDiffusionPipeline.from_pretrained(
        "CompVis/stable-diffusion-v1-4",
        torch_dtype=torch.float16
    ).to("cuda")
    
    # Load all embeddings
    for style_info in styles.values():
        embedding_path = hf_hub_download(
            repo_id=style_info["concept_url"],
            filename="learned_embeds.bin",
            repo_type="model"
        )
        pipe.load_textual_inversion(embedding_path)
    
    return pipe

def apply_purple_guidance(image, strength=0.5):
    """Apply purple guidance to an image"""
    img_array = np.array(image).astype(float)
    purple_mask = (img_array[:,:,0] > 100) & (img_array[:,:,2] > 100)
    img_array[purple_mask] = img_array[purple_mask] * (1 - strength) + np.array([128, 0, 128]) * strength
    return Image.fromarray(np.uint8(img_array.clip(0, 255)))

def generate_image(prompt, style, seed, apply_guidance, guidance_strength=0.5):
    """Generate an image with selected style and optional purple guidance"""
    if style not in styles:
        return None
    
    # Get style info
    style_info = styles[style]
    
    # Prepare generator
    generator = torch.Generator("cuda").manual_seed(int(seed))
    
    # Create styled prompt
    styled_prompt = f"{prompt} {style_info['token']}"
    
    # Generate image
    image = pipe(
        styled_prompt,
        generator=generator,
        guidance_scale=7.5,
        num_inference_steps=50
    ).images[0]
    
    # Apply purple guidance if requested
    if apply_guidance:
        image = apply_purple_guidance(image, guidance_strength)
    
    return image

# Initialize the pipeline globally
print("Loading pipeline and embeddings...")
pipe = load_pipeline()

# Create the Gradio interface
demo = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Textbox(label="Prompt", value="A serene mountain landscape with a lake at sunset"),
        gr.Radio(choices=list(styles.keys()), label="Style", value="glitch"),
        gr.Number(label="Seed", value=42),
        gr.Checkbox(label="Apply Purple Guidance", value=False),
        gr.Slider(minimum=0.0, maximum=1.0, value=0.5, label="Purple Guidance Strength")
    ],
    outputs=gr.Image(label="Generated Image"),
    title="Style-Guided Image Generation with Purple Enhancement",
    description="""Generate images in different styles with optional purple color guidance.
    Choose a style, enter a prompt, and optionally apply purple color enhancement.""",
    examples=[
        ["A serene mountain landscape with a lake at sunset", "glitch", 42, True, 0.5],
        ["A magical forest at twilight", "anime80s", 789, True, 0.7],
        ["A cyberpunk city at night", "night", 456, False, 0.5],
    ]
)

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