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()