app.py

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

import subprocess
subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)

# Specify the path to the model
model_path = "deepseek-ai/Janus-1.3B"

# Load the VLChatProcessor and tokenizer
print("Loading VLChatProcessor and tokenizer...")
vl_chat_processor: VLChatProcessor = VLChatProcessor.from_pretrained(model_path)
tokenizer = vl_chat_processor.tokenizer

# Load the MultiModalityCausalLM model
print("Loading MultiModalityCausalLM model...")
vl_gpt: MultiModalityCausalLM = AutoModelForCausalLM.from_pretrained(
    model_path, trust_remote_code=True
)
# Move the model to GPU with bfloat16 precision for efficiency
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
vl_gpt = vl_gpt.to(torch.bfloat16 if device.type == "cuda" else torch.float32).to(device).eval()

@spaces.GPU(duration=120)
def text_image_to_text(user_text: str, user_image: Image.Image) -> str:
    """
    Generate a textual response based on user-provided text and image.
    This can be used for tasks like converting an image of a formula to LaTeX code
    or generating descriptive captions.
    """
    # Define the conversation with user-provided text and image
    conversation = [
        {
            "role": "User",
            "content": user_text,
            "images": [user_image],
        },
        {"role": "Assistant", "content": ""},
    ]
    
    # Load the PIL images from the conversation
    pil_images = load_pil_images(conversation)
    
    # Prepare the inputs for the model
    prepare_inputs = vl_chat_processor(
        conversations=conversation, images=pil_images, force_batchify=True
    ).to(device)
    
    # Prepare input embeddings
    inputs_embeds = vl_gpt.prepare_inputs_embeds(**prepare_inputs)
    
    # Generate the response from the model
    with torch.no_grad():
        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,
            use_cache=True,
        )
    
    # Decode the generated tokens to get the answer
    answer = tokenizer.decode(outputs[0].cpu().tolist(), skip_special_tokens=True)
    
    return answer

@spaces.GPU(duration=120)
def text_to_image(prompt: str) -> Image.Image:
    """
    Generate an image based on the input text prompt.
    """
    # Define the conversation with the user prompt
    conversation = [
        {
            "role": "User",
            "content": prompt,
        },
        {"role": "Assistant", "content": ""},
    ]
    
    # Apply the SFT template to format the prompt
    sft_format = vl_chat_processor.apply_sft_template_for_multi_turn_prompts(
        conversations=conversation,
        sft_format=vl_chat_processor.sft_format,
        system_prompt="",
    )
    prompt_text = sft_format + vl_chat_processor.image_start_tag
    
    # Encode the prompt
    input_ids = vl_chat_processor.tokenizer.encode(prompt_text)
    input_ids = torch.LongTensor(input_ids).unsqueeze(0).to(device)
    
    # Prepare tokens for generation
    parallel_size = 1  # Adjust based on GPU memory
    tokens = torch.zeros((parallel_size*2, len(input_ids[0])), dtype=torch.int).to(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
    
    # Get input embeddings
    inputs_embeds = vl_gpt.language_model.get_input_embeddings()(tokens)
    
    # Generation parameters
    image_token_num_per_image = 576
    img_size = 384
    patch_size = 16
    cfg_weight = 5
    temperature = 1
    
    # Initialize tensor to store generated tokens
    generated_tokens = torch.zeros((parallel_size, image_token_num_per_image), dtype=torch.int).to(device)
    
    for i in range(image_token_num_per_image):
        if i == 0:
            outputs = vl_gpt.language_model.model(inputs_embeds=inputs_embeds, use_cache=True)
        else:
            outputs = vl_gpt.language_model.model(inputs_embeds=inputs_embeds, use_cache=True, past_key_values=outputs.past_key_values)
        
        hidden_states = outputs.last_hidden_state
        
        # Get logits and apply classifier-free guidance
        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)
        
        # Sample the next token
        probs = torch.softmax(logits / temperature, dim=-1)
        next_token = torch.multinomial(probs, num_samples=1)
        generated_tokens[:, i] = next_token.squeeze(dim=-1)
        
        # Prepare for the next step
        next_token_combined = torch.cat([next_token, next_token], dim=0).view(-1)
        img_embeds = vl_gpt.prepare_gen_img_embeds(next_token_combined)
        inputs_embeds = img_embeds.unsqueeze(dim=1)
    
    # Decode the generated tokens to get the image
    dec = vl_gpt.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)
    
    # Convert to PIL Image
    visual_img = dec[0]
    image = Image.fromarray(visual_img)
    
    return image

# Create the Gradio interface
with gr.Blocks() as demo:
    gr.Markdown(
        """
        # Janus-1.3B Gradio Demo
        This demo showcases two functionalities using the Janus-1.3B model:
        1. **Text + Image to Text**: Input both text and an image to generate a textual response.
        2. **Text to Image**: Enter a descriptive text prompt to generate a corresponding image.
        """
    )
    
    with gr.Tab("Text + Image to Text"):
        gr.Markdown("### Generate Text Based on Input Text and Image")
        with gr.Row():
            with gr.Column():
                user_text_input = gr.Textbox(
                    lines=2,
                    placeholder="Enter your instructions or description here...",
                    label="Input Text",
                )
                user_image_input = gr.Image(
                    type="pil",
                    label="Upload Image",
                )
                submit_btn = gr.Button("Generate Text")
            with gr.Column():
                text_output = gr.Textbox(
                    label="Generated Text",
                    lines=15,
                    interactive=False,
                )
        submit_btn.click(fn=text_image_to_text, inputs=[user_text_input, user_image_input], outputs=text_output)
    
    with gr.Tab("Text to Image"):
        gr.Markdown("### Generate Image Based on Text Prompt")
        with gr.Row():
            with gr.Column():
                prompt_input = gr.Textbox(
                    lines=2,
                    placeholder="Enter your image description here...",
                    label="Text Prompt",
                )
                generate_btn = gr.Button("Generate Image")
            with gr.Column():
                image_output = gr.Image(
                    label="Generated Image",
                )
        generate_btn.click(fn=text_to_image, inputs=prompt_input, outputs=image_output)
    
# Launch the Gradio app
if __name__ == "__main__":
    demo.launch()