app.py

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, seed

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

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

demo.launch()