app.py

import gradio as gr
import torch
from src.transformer import SymmetricTransformer2DModel
from src.pipeline import UnifiedPipeline
from src.scheduler import Scheduler
from torchvision import transforms
from transformers import CLIPTextModelWithProjection, CLIPTokenizer
from diffusers import VQModel
import os
from PIL import Image
import numpy as np
import spaces 

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


def load_models(model_path="MeissonFlow/Meissonic", 
            transformer_path="MeissonFlow/Muddit"):
    model = SymmetricTransformer2DModel.from_pretrained(
        transformer_path,
        subfolder="1024/transformer"
    )
    vq_model = VQModel.from_pretrained(model_path, subfolder="vqvae")
    text_encoder = CLIPTextModelWithProjection.from_pretrained(model_path, subfolder="text_encoder")
    tokenizer = CLIPTokenizer.from_pretrained(model_path, subfolder="tokenizer")
    scheduler = Scheduler.from_pretrained(model_path, subfolder="scheduler")

    pipe = UnifiedPipeline(
        vqvae=vq_model,
        tokenizer=tokenizer,
        text_encoder=text_encoder,
        transformer=model,
        scheduler=scheduler,
    )
    return pipe

# Load models (global variable to avoid reloading)
pipe = load_models()
pipe.to(device)

# Common transform
def get_transform(resolution):
    return transforms.Compose([
        transforms.Resize((resolution, resolution)),
        transforms.ToTensor(),
    ])

# Image-to-Text Function
@spaces.GPU
def image_to_text(image, prompt, seed=42, steps=64, cfg=9.0):
    try:
        resolution = 1024
        transform = get_transform(resolution)
        
        if image is not None:
            pil_image = Image.fromarray(image.astype('uint8'), 'RGB') if isinstance(image, np.ndarray) else image
            images = torch.stack([transform(pil_image)])
            questions = [prompt] if prompt else ["Please describe this image."]
        else:
            images = None
            questions = [prompt] if prompt else ["Please generate an image description."]

        output = pipe(
            prompt=questions,
            image=images,
            height=resolution,
            width=resolution,
            guidance_scale=cfg,
            num_inference_steps=steps,
            mask_token_embedding="./mask_token_embedding.pth",
            generator=torch.manual_seed(seed),
        )

        return output.prompts[0]
    
    except Exception as e:
        return f"Error: {str(e)}"

# Text-to-Image Function
@spaces.GPU
def text_to_image(prompt, negative_prompt, num_images=1, seed=42, steps=64, cfg=9.0):
    try:
        resolution = 1024
        negative_prompt = negative_prompt or "worst quality, low quality, low res, blurry, distortion, watermark, logo, signature, text, jpeg artifacts, signature, sketch, duplicate, ugly, identifying mark"
        
        output = pipe(
            prompt=[prompt]*num_images,
            negative_prompt=[negative_prompt]*num_images,
            height=resolution,
            width=resolution,
            guidance_scale=cfg,
            num_inference_steps=steps,
            mask_token_embedding="./mask_token_embedding.pth",
            generator=torch.manual_seed(seed),
        )
        
        return output.images
    
    except Exception as e:
        print(f"Error: {str(e)}")
        return None

# Create Gradio interface with Soft theme
with gr.Blocks(theme=gr.themes.Soft(), title="Muddit Unifined Model") as demo:
    gr.Markdown("# 🌌 Muddit: Liberating Generation Beyond Text-to-Image with a Unified Discrete Diffusion Model.")
    gr.Markdown(" Muddit is a unified discrete diffusion transformer that enables fast and parallel generation across both text and image modalities.")
    
    with gr.Tab("Image to Text"):
        with gr.Row():
            with gr.Column():
                i2t_image_input = gr.Image(label="Upload Image", type="pil")
                i2t_prompt_input = gr.Textbox(label="Prompt", value="Please describe this image.", placeholder="Enter your prompt here...")
                
                with gr.Accordion("Advanced Settings", open=False):
                    seed = gr.Slider(label="Seed", minimum=0, maximum=2**32 - 1, step=1, value=42)
                    i2t_steps = gr.Slider(label="Inference Steps", minimum=10, maximum=100, value=64, step=1)
                    i2t_cfg = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=20.0, value=9.0, step=0.5)
                
                i2t_submit_btn = gr.Button("Generate Description", variant="primary")
            
            with gr.Column():
                i2t_output_text = gr.Textbox(label="Generated Description", interactive=False)
                i2t_examples = gr.Examples(
                    examples=[
                        ["assets/man.jpg"],
                        ["assets/tennis.jpg"],
                        ["assets/pizza2.jpg"],
                        ["assets/plane.jpg"],
                        ["assets/zebra.jpg"],
                        ["assets/building.jpg"],
                        ["assets/flower.jpg"],
                    ],
                    inputs=[i2t_image_input],
                    label="Example Inputs"
                )
    
    with gr.Tab("VQA"):
        with gr.Row():
            with gr.Column():
                vqa_image_input = gr.Image(label="Upload Image", type="pil")
                vqa_prompt_input = gr.Textbox(label="Prompt", placeholder="Enter your question here...")
                
                with gr.Accordion("Advanced Settings", open=False):
                    seed = gr.Slider(label="Seed", minimum=0, maximum=2**32 - 1, step=1, value=42)
                    vqa_steps = gr.Slider(label="Inference Steps", minimum=10, maximum=100, value=64, step=1)
                    vqa_cfg = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=20.0, value=9.0, step=0.5)
                
                vqa_submit_btn = gr.Button("Generate Answer", variant="primary")
            
            with gr.Column():
                vqa_output_text = gr.Textbox(label="Generated Answer", interactive=False)
                vqa_examples = gr.Examples(
                    examples=[
                        ["assets/kid.jpg", "What color is the kid's hair?"],
                        ["assets/street.jpg", "Can someone legally walk across the street right now?"],
                        ["assets/dog.jpg", "Where is the dog laying?"],
                        ["assets/dog2.jpg", "What color is the toy the dog is holding?"],
                        ["assets/pizza.jpg", "What food item is shown?"],
                        ["assets/sheep.jpg", "How many sheep are pictured?"],
                        ["assets/car.jpg", "Where are the cars?"],
                    ],
                    inputs=[vqa_image_input, vqa_prompt_input],
                    label="Example Inputs"
                )
    
    with gr.Tab("Text to Image"):
        with gr.Row():
            with gr.Column():
                t2i_prompt_input = gr.Textbox(label="Prompt", placeholder="Describe the image you want to generate...")
                t2i_negative_prompt = gr.Textbox(label="Negative Prompt", 
                                               value="worst quality, low quality, low res, blurry, distortion, watermark, logo, signature, text, jpeg artifacts, signature, sketch, duplicate, ugly, identifying mark",
                                               placeholder="What you don't want in the image...",
                                               lines=5)
                t2i_num_images = gr.Slider(label="Number of Images", minimum=1, maximum=4, value=1, step=1)
                
                with gr.Accordion("Advanced Settings", open=False):
                    seed = gr.Slider(label="Seed", minimum=0, maximum=2**32 - 1, step=1, value=42)
                    t2i_steps = gr.Slider(label="Inference Steps", minimum=10, maximum=100, value=64, step=1)
                    t2i_cfg = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=20.0, value=9.0, step=0.5)
                
                t2i_submit_btn = gr.Button("Generate Images", variant="primary")
            
            with gr.Column():
                t2i_gallery = gr.Gallery(label="Generated Images")
                t2i_examples = gr.Examples(
                    examples=[
                        ["A line art portrait showcasing a human figure with flowing, textured strokes"],
                        ["A hyper realistic image of a chimpanzee with a glass-enclosed brain on his head, standing amidst lush, bioluminescent foliage in a vibrant futuristic forest"],
                        ["A samurai in a stylized cyberpunk outfit adorned with intricate steampunk gear and floral accents, his Mandalorian armor gleaming under the backlighting"],
                        ["A translucent, minimalist Porsche 911 GT3RS built from sleek carbon fiber, its aerodynamic body designed in the spirit of '60s Braun and modern Apple minimalism"],
                        ["A realistic photograph of a ramadan tent shaped like a crescent moon under a velvety back sky studded with the milky way"],
                        ["A portrait of John Lennon, captured in the gritty detail of line art"],
                        ["In a world plunged into an unending darkness, remnants of fading starlight pierce through a heavy, smog-filled sky"]
                    ],
                    inputs=[t2i_prompt_input],
                    label="Example Prompts"
                )
    
    # Event handlers
    i2t_submit_btn.click(
        fn=image_to_text,
        inputs=[i2t_image_input, i2t_prompt_input, seed, i2t_steps, i2t_cfg],
        outputs=i2t_output_text
    )
    
    vqa_submit_btn.click(
        fn=image_to_text,
        inputs=[vqa_image_input, vqa_prompt_input, seed, vqa_steps, vqa_cfg],
        outputs=vqa_output_text
    )
    
    t2i_submit_btn.click(
        fn=text_to_image,
        inputs=[t2i_prompt_input, t2i_negative_prompt, t2i_num_images, seed, t2i_steps, t2i_cfg],
        outputs=t2i_gallery
    )

demo.launch()