import os
import sys
sys.path.append("./")

import torch
from torchvision import transforms
from src.transformer import Transformer2DModel
from src.pipeline import Pipeline
from src.scheduler import Scheduler
from transformers import (
    CLIPTextModelWithProjection,
    CLIPTokenizer,
)
from diffusers import VQModel
import gradio as gr
import spaces 

device = 'cuda' if torch.cuda.is_available() else 'cpu'
dtype = torch.bfloat16

model_path = "Collov-Labs/Monetico"

model = Transformer2DModel.from_pretrained(model_path, subfolder="transformer", torch_dtype=dtype)
vq_model = VQModel.from_pretrained(model_path, subfolder="vqvae", torch_dtype=dtype)
text_encoder = CLIPTextModelWithProjection.from_pretrained(
            "laion/CLIP-ViT-H-14-laion2B-s32B-b79K", torch_dtype=dtype
        )
tokenizer = CLIPTokenizer.from_pretrained(model_path, subfolder="tokenizer", torch_dtype=dtype)
scheduler = Scheduler.from_pretrained(model_path, subfolder="scheduler", torch_dtype=dtype)
pipe = Pipeline(vq_model, tokenizer=tokenizer, text_encoder=text_encoder, transformer=model, scheduler=scheduler)
pipe.to(device)

MAX_SEED = 2**32 - 1

@spaces.GPU
def generate_image(occasion, theme, colors, randomize_seed=True, seed=0):
    prompt = f"{occasion} theme: {theme}, colors: {colors} design inspiration"
    if randomize_seed or seed == 0:
        seed = torch.randint(0, MAX_SEED, (1,)).item()
    torch.manual_seed(seed)
    
    image = pipe(
        prompt=prompt,
        height=512,
        width=512,
        guidance_scale=9.0,
        num_inference_steps=50
    ).images[0]
    
    return image 

css = """
#col-container {
    margin: 0 auto;
    max-width: 640px;
}
"""

examples = [
    ["Corporate Anniversary", "Legacy & Growth", "navy and silver"],
    ["Product Launch", "Innovation Spark", "blue and white"],
    ["Team Appreciation", "Together We Thrive", "green and gold"],
    ["Award Ceremony", "Excellence Awards", "black and gold"],
    ["Milestone Celebration", "10 Years Strong", "emerald green and silver"],
    ["Holiday Party", "Winter Festivity", "silver and blue"],
    ["Sales Achievement", "Peak Performers", "crimson and gray"],
    ["Client Appreciation", "Thank You Event", "ivory and gold"],
    ["Office Opening", "New Beginnings", "teal and white"],
    ["Retirement Celebration", "Years of Dedication", "bronze and navy"],
    ["Quarterly Town Hall", "United Vision", "purple and silver"],
    ["Annual Conference", "Forward Together", "black and royal blue"],
    ["Workshop Event", "Skill Building", "orange and gray"],
    ["Networking Gala", "Professional Connections", "champagne and gold"],
    ["Leadership Retreat", "Inspire & Lead", "forest green and white"],
]


with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown("# Cake & Gift Design Inspiration")
        with gr.Row():
            occasion = gr.Text(label="Occasion", placeholder="Enter occasion, e.g., Wedding, Birthday")
            theme = gr.Text(label="Theme", placeholder="Enter theme, e.g., Vintage, Space Adventure")
            colors = gr.Text(label="Colors", placeholder="Enter colors, e.g., white and gold")
            run_button = gr.Button("Generate Design", variant="primary")
        result = gr.Image(label="Generated Design", show_label=False)
        gr.Examples(examples=examples, inputs=[occasion, theme, colors])
        
    gr.on(
        triggers=[run_button.click],
        fn=generate_image,
        inputs=[occasion, theme, colors],
        outputs=[result],  # Expect only the image output
    )

demo.launch()