import os
import torch
import numpy as np
from PIL import Image
import gradio as gr
from transformers import AutoModelForCausalLM
from janus.models import MultiModalityCausalLM, VLChatProcessor
import spaces

# === Load model and processor ===
model_path = "FreedomIntelligence/Janus-4o-7B"
vl_chat_processor: VLChatProcessor = VLChatProcessor.from_pretrained(model_path)
tokenizer = vl_chat_processor.tokenizer

vl_gpt: MultiModalityCausalLM = AutoModelForCausalLM.from_pretrained(
    model_path, trust_remote_code=True, torch_dtype=torch.bfloat16
)
vl_gpt = vl_gpt.cuda().eval()

# === Image generation function ===
def text_to_image_generate(input_prompt, output_path, vl_chat_processor, vl_gpt, temperature=1.0, parallel_size=1, cfg_weight=5):
    torch.cuda.empty_cache()

    # Apply prompt formatting
    conversation = [
        {"role": "<|User|>", "content": input_prompt},
        {"role": "<|Assistant|>", "content": ""},
    ]

    sft_format = vl_chat_processor.apply_sft_template_for_multi_turn_prompts(
        conversations=conversation,
        sft_format=vl_chat_processor.sft_format,
        system_prompt="",
    )

    prompt = sft_format + vl_chat_processor.image_start_tag
    mmgpt = vl_gpt

    image_token_num = 576
    img_size = 384
    patch_size = 16

    with torch.inference_mode():
        input_ids = tokenizer.encode(prompt)
        input_ids = torch.LongTensor(input_ids)

        tokens = torch.zeros((parallel_size * 2, len(input_ids)), dtype=torch.int).cuda()
        for i in range(parallel_size * 2):
            tokens[i, :] = input_ids
            if i % 2 != 0:
                tokens[i, 1:-1] = tokenizer.pad_token_id  # More robust

        inputs_embeds = mmgpt.language_model.get_input_embeddings()(tokens)
        generated_tokens = torch.zeros((parallel_size, image_token_num), dtype=torch.int).cuda()

        for i in range(image_token_num):
            outputs = mmgpt.language_model.model(inputs_embeds=inputs_embeds, use_cache=False)
            hidden_states = outputs.last_hidden_state
            logits = mmgpt.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, next_token], dim=1).reshape(-1)
            img_embeds = mmgpt.prepare_gen_img_embeds(next_token)
            inputs_embeds = img_embeds.unsqueeze(1)

        # Decode image
        dec = mmgpt.gen_vision_model.decode_code(
            generated_tokens.to(dtype=torch.int),
            shape=[parallel_size, 8, img_size // patch_size, img_size // patch_size]
        )
        dec = dec.to(torch.float32).cpu().numpy().transpose(0, 2, 3, 1)
        dec = np.clip((dec + 1) / 2 * 255, 0, 255).astype(np.uint8)

        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        output_images = []
        for i in range(parallel_size):
            save_path = output_path.replace('.png', f'_{i}.png')
            Image.fromarray(dec[i]).save(save_path)
            output_images.append(save_path)

        return output_images

# === Gradio handler ===
@spaces.GPU(duration=120)
def janus_generate_image(message, history):
    prompt = message
    output_path = "./output/image.png"
    images = text_to_image_generate(prompt, output_path, vl_chat_processor, vl_gpt, parallel_size=1)
    return {"role": "assistant", "content": images[0]}

# === Gradio UI ===
demo = gr.ChatInterface(
    fn=janus_generate_image,
    title="Janus Text-to-Image",
    description="Generate images from natural language prompts using Janus-4o-7B",
    additional_inputs=[],
    chatbot=gr.Chatbot(show_copy_button=True),
    examples=["a cat", "a spaceship landing on Mars", "a fantasy castle at sunset"],
    theme="soft",
)

if __name__ == "__main__":
    demo.launch()