import gradio as gr
import torch
from transformers import AutoConfig, AutoModelForCausalLM
from janus.models import MultiModalityCausalLM, VLChatProcessor
from janus.utils.io import load_pil_images
from PIL import Image
import numpy as np
import os
import time
from Upsample import RealESRGAN
import spaces  # Import spaces for ZeroGPU compatibility

# ---------------------------
# Load model and processor
# ---------------------------
model_path = "deepseek-ai/Janus-Pro-7B"
config = AutoConfig.from_pretrained(model_path)
language_config = config.language_config
language_config._attn_implementation = 'eager'
vl_gpt = AutoModelForCausalLM.from_pretrained(model_path,
                                             language_config=language_config,
                                             trust_remote_code=True)
if torch.cuda.is_available():
    vl_gpt = vl_gpt.to(torch.bfloat16).cuda()
else:
    vl_gpt = vl_gpt.to(torch.float16)

vl_chat_processor = VLChatProcessor.from_pretrained(model_path)
tokenizer = vl_chat_processor.tokenizer
cuda_device = 'cuda' if torch.cuda.is_available() else 'cpu'

# SR (Super Resolution) model
sr_model = RealESRGAN(torch.device('cuda' if torch.cuda.is_available() else 'cpu'), scale=2)
sr_model.load_weights(f'weights/RealESRGAN_x2.pth', download=False)

# ---------------------------
# Multimodal Understanding Function
# ---------------------------
@torch.inference_mode()
@spaces.GPU(duration=120)
def multimodal_understanding(image, question, seed, top_p, temperature, progress=gr.Progress(track_tqdm=True)):
    # Clear CUDA cache before generating
    torch.cuda.empty_cache()

    # Set seed for reproducibility
    torch.manual_seed(seed)
    np.random.seed(seed)
    torch.cuda.manual_seed(seed)
    
    # Prepare conversation – note the use of a placeholder for the image.
    conversation = [
        {
            "role": "<|User|>",
            "content": f"<image_placeholder>\n{question}",
            "images": [image],
        },
        {"role": "<|Assistant|>", "content": ""},
    ]
    
    # The chat processor expects PIL images.
    pil_images = [Image.fromarray(np.array(image))] if not isinstance(image, Image.Image) else [image]
    prepare_inputs = vl_chat_processor(
        conversations=conversation, images=pil_images, force_batchify=True
    ).to(cuda_device, dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float16)
    
    inputs_embeds = vl_gpt.prepare_inputs_embeds(**prepare_inputs)
    
    outputs = vl_gpt.language_model.generate(
        inputs_embeds=inputs_embeds,
        attention_mask=prepare_inputs.attention_mask,
        pad_token_id=tokenizer.eos_token_id,
        bos_token_id=tokenizer.bos_token_id,
        eos_token_id=tokenizer.eos_token_id,
        max_new_tokens=512,
        do_sample=False if temperature == 0 else True,
        use_cache=True,
        temperature=temperature,
        top_p=top_p,
    )
    
    answer = tokenizer.decode(outputs[0].cpu().tolist(), skip_special_tokens=True)
    return answer

# ---------------------------
# Image Generation Functions
# ---------------------------
def generate(input_ids,
             width,
             height,
             temperature: float = 1,
             parallel_size: int = 5,
             cfg_weight: float = 5,
             image_token_num_per_image: int = 576,
             patch_size: int = 16,
             progress=gr.Progress(track_tqdm=True)):
    torch.cuda.empty_cache()
    
    tokens = torch.zeros((parallel_size * 2, len(input_ids)), dtype=torch.int).to(cuda_device)
    for i in range(parallel_size * 2):
        tokens[i, :] = input_ids
        if i % 2 != 0:
            tokens[i, 1:-1] = vl_chat_processor.pad_id
    inputs_embeds = vl_gpt.language_model.get_input_embeddings()(tokens)
    generated_tokens = torch.zeros((parallel_size, image_token_num_per_image), dtype=torch.int).to(cuda_device)

    pkv = None
    for i in range(image_token_num_per_image):
        with torch.no_grad():
            outputs = vl_gpt.language_model.model(inputs_embeds=inputs_embeds,
                                                  use_cache=True,
                                                  past_key_values=pkv)
            pkv = outputs.past_key_values
            hidden_states = outputs.last_hidden_state
            logits = vl_gpt.gen_head(hidden_states[:, -1, :])
            logit_cond = logits[0::2, :]
            logit_uncond = logits[1::2, :]
            logits = logit_uncond + cfg_weight * (logit_cond - logit_uncond)
            probs = torch.softmax(logits / temperature, dim=-1)
            next_token = torch.multinomial(probs, num_samples=1)
            generated_tokens[:, i] = next_token.squeeze(dim=-1)
            next_token = torch.cat([next_token.unsqueeze(dim=1), next_token.unsqueeze(dim=1)], dim=1).view(-1)

            img_embeds = vl_gpt.prepare_gen_img_embeds(next_token)
            inputs_embeds = img_embeds.unsqueeze(dim=1)
    
    patches = vl_gpt.gen_vision_model.decode_code(generated_tokens.to(dtype=torch.int),
                                                  shape=[parallel_size, 8, width // patch_size, height // patch_size])
    return generated_tokens.to(dtype=torch.int), patches

def unpack(dec, width, height, parallel_size=5):
    dec = dec.to(torch.float32).cpu().numpy().transpose(0, 2, 3, 1)
    dec = np.clip((dec + 1) / 2 * 255, 0, 255)
    visual_img = np.zeros((parallel_size, width, height, 3), dtype=np.uint8)
    visual_img[:, :, :] = dec
    return visual_img

@torch.inference_mode()
@spaces.GPU(duration=120)
def generate_image(prompt,
                   seed=None,
                   guidance=5,
                   t2i_temperature=1.0,
                   progress=gr.Progress(track_tqdm=True)):
    torch.cuda.empty_cache()
    if seed is not None:
        torch.manual_seed(seed)
        torch.cuda.manual_seed(seed)
        np.random.seed(seed)
    width = 384
    height = 384
    parallel_size = 4

    with torch.no_grad():
        messages = [{'role': '<|User|>', 'content': prompt},
                    {'role': '<|Assistant|>', 'content': ''}]
        text = vl_chat_processor.apply_sft_template_for_multi_turn_prompts(conversations=messages,
                                                                           sft_format=vl_chat_processor.sft_format,
                                                                           system_prompt='')
        text = text + vl_chat_processor.image_start_tag
        
        input_ids = torch.LongTensor(tokenizer.encode(text))
        output, patches = generate(input_ids,
                                   width // 16 * 16,
                                   height // 16 * 16,
                                   cfg_weight=guidance,
                                   parallel_size=parallel_size,
                                   temperature=t2i_temperature)
        images = unpack(patches,
                        width // 16 * 16,
                        height // 16 * 16,
                        parallel_size=parallel_size)
        # Upsample the generated images
        stime = time.time()
        ret_images = [image_upsample(Image.fromarray(images[i])) for i in range(parallel_size)]
        print(f'upsample time: {time.time() - stime}')
        return ret_images  # returns a list

@spaces.GPU(duration=60)
def image_upsample(img: Image.Image) -> Image.Image:
    if img is None:
        raise Exception("Image not uploaded")
    
    width, height = img.size
    if width >= 5000 or height >= 5000:
        raise Exception("The image is too large.")

    global sr_model
    result = sr_model.predict(img.convert('RGB'))
    return result

# A helper function to generate a single image (the first result) from a description.
def generate_single_image(prompt, seed, guidance, t2i_temperature):
    images = generate_image(prompt, seed, guidance, t2i_temperature)
    # Return the first image (if available)
    return images[0] if images else None

# ---------------------------
# Chat About Generated Image
# ---------------------------
# This function uses the generated image and a chat question.
def chat_about_image(generated_image, chat_text, seed, top_p, temperature, chat_history):
    if generated_image is None:
        return chat_history, "Please generate an image first by entering a description above."
    response = multimodal_understanding(generated_image, chat_text, seed, top_p, temperature)
    chat_history.append((chat_text, response))
    return chat_history, ""

# ---------------------------
# Gradio Interface
# ---------------------------
css = '''
.gradio-container {max-width: 960px !important}
'''

with gr.Blocks(css=css, title="Janus Pro 7B – Image Generation and Chat") as demo:
    gr.Markdown("# Janus Pro 7B: Image Generation and Conversation")
    gr.Markdown("Enter an image description below to have the model generate an image. Once generated, you can chat about the image and ask questions.")
    
    # States to store the generated image and the chat history.
    state_image = gr.State(None)
    state_history = gr.State([])

    with gr.Row():
        with gr.Column():
            gr.Markdown("### Step 1. Generate an Image from Description")
            description_input = gr.Textbox(label="Image Description", placeholder="Describe the image you want...")
            with gr.Accordion("Advanced Generation Options", open=False):
                gen_seed_input = gr.Number(label="Seed", precision=0, value=42)
                guidance_input = gr.Slider(minimum=1, maximum=10, value=5, step=0.5, label="CFG Weight")
                t2i_temperature_input = gr.Slider(minimum=0, maximum=1, value=1.0, step=0.05, label="Temperature")
            generate_button = gr.Button("Generate Image")
            image_output = gr.Image(label="Generated Image", interactive=False)
        with gr.Column():
            gr.Markdown("### Step 2. Chat about the Image")
            gr.Markdown("Ask questions or discuss the generated image below. (If no image has been generated yet, please do so in Step 1.)")
            with gr.Accordion("Advanced Chat Options", open=False):
                chat_seed_input = gr.Number(label="Seed", precision=0, value=42)
                top_p_input = gr.Slider(minimum=0, maximum=1, value=0.95, step=0.05, label="top_p")
                chat_temperature_input = gr.Slider(minimum=0, maximum=1, value=0.1, step=0.05, label="Temperature")
            chatbox = gr.Chatbot(label="Conversation")
            chat_input = gr.Textbox(label="Your Message", placeholder="Enter your question or comment here...")
            send_button = gr.Button("Send")
    
    # When the user clicks the "Generate Image" button:
    generate_button.click(
        fn=generate_single_image,
        inputs=[description_input, gen_seed_input, guidance_input, t2i_temperature_input],
        outputs=image_output
    ).then(
        fn=lambda img: img,  # pass through the generated image
        inputs=image_output,
        outputs=state_image
    )
    
    # When the user sends a chat message, update the conversation.
    send_button.click(
        fn=chat_about_image,
        inputs=[state_image, chat_input, chat_seed_input, top_p_input, chat_temperature_input, state_history],
        outputs=[chatbox, chat_input],
    )

demo.launch(share=True)