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'

model_path = "MeissonFlow/Meissonic"
model = Transformer2DModel.from_pretrained(model_path, subfolder="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 = Pipeline(vq_model, tokenizer=tokenizer, text_encoder=text_encoder, transformer=model, scheduler=scheduler)
pipe.to(device)

@spaces.GPU
def generate_image(prompt, negative_prompt, resolution, steps, cfg, progress=gr.Progress(track_tqdm=True)):
    image = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        height=resolution,
        width=resolution,
        guidance_scale=cfg,
        num_inference_steps=steps
    ).images[0]
    
    return image

# Default negative prompt
default_negative_prompt = "worst quality, normal quality, low quality, low res, blurry, distortion, text, watermark, logo, banner, extra digits, cropped, jpeg artifacts, signature, username, error, sketch, duplicate, ugly, monochrome, horror, geometry, mutation, disgusting, bad anatomy, bad proportions, bad quality, deformed, disconnected limbs, out of frame, out of focus, dehydrated, disfigured, extra arms, extra limbs, extra hands, fused fingers, gross proportions, long neck, jpeg, malformed limbs, mutated, mutated hands, mutated limbs, missing arms, missing fingers, picture frame, poorly drawn hands, poorly drawn face, collage, pixel, pixelated, grainy, color aberration, amputee, autograph, bad illustration, beyond the borders, blank background, body out of frame, boring background, branding, cut off, dismembered, disproportioned, distorted, draft, duplicated features, extra fingers, extra legs, fault, flaw, grains, hazy, identifying mark, improper scale, incorrect physiology, incorrect ratio, indistinct, kitsch, low resolution"

# Gradio interface
iface = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Textbox(label="Prompt", placeholder="Enter your prompt here..."),
        gr.Textbox(label="Negative Prompt", value=default_negative_prompt),
        gr.Slider(512, 1024, 1024, step=64, label="Resolution"),
        gr.Slider(1, 100, 48, step=1, label="Number of Steps"),
        gr.Slider(1, 20, 9, step=0.5, label="CFG Scale")
    ],
    outputs=gr.Image(type="pil"),
    title="Meissonic Image Generator",
    description="Generate images using the Meissonic model. Enter a prompt and adjust parameters to create your image.",
)

iface.launch()