first commit

app.py

@@ -1,7 +1,16 @@
-import gradio as gr
+import torch
+import spaces
+from model_loader import load_model_and_processor
+from image_generator import process_and_generate
+from gradio_interface import create_gradio_interface
 
-def greet(name):
-    return "Hello " + name + "!!"
+if __name__ == "__main__":
+    # Set the model path
+    model_path = "deepseek-ai/Janus-1.3B"
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+    # Load the model and processor
+    vl_gpt, vl_chat_processor = load_model_and_processor(model_path)
+
+    # Create and launch the Gradio interface
+    demo = create_gradio_interface(vl_gpt, vl_chat_processor, process_and_generate)
+    demo.launch(allowed_paths=["/"])

gradio_interface.py

@@ -0,0 +1,32 @@
+import gradio as gr
+from PIL import Image
+
+def create_gradio_interface(vl_gpt, vl_chat_processor, process_and_generate):
+    def gradio_process_and_generate(input_image, prompt, num_images, cfg_weight):
+        return process_and_generate(vl_gpt, vl_chat_processor, input_image, prompt, num_images, cfg_weight)
+
+    explanation = """Janus 1.3B uses a differerent visual encoder for understanding and generation.
+
+![Janus Model Architecture](file/images/janus_architecture.svg)
+
+Here, by feeding the model an image and then asking it to generate that same image, we visualize the model's ability to translate input (understanding) embedding space to generative embedding space."""
+
+    with gr.Blocks() as demo:
+        gr.Markdown("# How Janus-1.3B sees itself")
+
+        with gr.Row():
+            input_image = gr.Image(type="filepath", label="Input Image")
+            output_images = gr.Gallery(label="Generated Images", columns=2, rows=2)
+        prompt = gr.Textbox(label="Prompt", value="Exactly what is shown in the image.")
+        num_images = gr.Slider(minimum=1, maximum=12, value=12, step=1, label="Number of Images to Generate")
+        cfg_weight = gr.Slider(minimum=1, maximum=10, value=5, step=0.1, label="CFG Weight")
+        generate_btn = gr.Button("Generate", variant="primary", size="lg")
+        
+        generate_btn.click(
+            fn=gradio_process_and_generate,
+            inputs=[input_image, prompt, num_images, cfg_weight],
+            outputs=output_images
+        )
+        gr.Markdown(explanation)
+
+    return demo

image_generator.py

@@ -0,0 +1,117 @@
+import os
+import PIL.Image
+import torch
+import numpy as np
+from janus.utils.io import load_pil_images
+from janus.models import MultiModalityCausalLM, VLChatProcessor
+from functools import lru_cache
+
+
+def prepare_classifier_free_guidance_input(input_embeds, vl_chat_processor, mmgpt, batch_size=16):
+    uncond_input_ids = torch.full((1, input_embeds.shape[1]), 
+                                  vl_chat_processor.pad_id, 
+                                  dtype=torch.long, 
+                                  device=input_embeds.device)
+    uncond_input_ids[:, 0] = input_embeds.shape[1] - 1
+    uncond_input_ids[:, -1] = vl_chat_processor.tokenizer.eos_token_id
+    
+    uncond_input_embeds = mmgpt.language_model.get_input_embeddings()(uncond_input_ids)
+    uncond_input_embeds[:, -1, :] = input_embeds[:, -1, :]
+
+    cond_input_embeds = input_embeds.repeat(batch_size, 1, 1)
+    uncond_input_embeds = uncond_input_embeds.repeat(batch_size, 1, 1)
+    
+    combined_input_embeds = torch.stack([cond_input_embeds, uncond_input_embeds], dim=1)
+    combined_input_embeds = combined_input_embeds.view(batch_size * 2, -1, input_embeds.shape[-1])
+    
+    return combined_input_embeds
+
+@spaces.GPU
+@torch.inference_mode()
+def generate(
+    mmgpt: MultiModalityCausalLM,
+    vl_chat_processor: VLChatProcessor,
+    inputs_embeds,
+    temperature: float = 1,
+    parallel_size: int = 1,
+    cfg_weight: float = 5,
+    image_token_num_per_image: int = 576,
+    img_size: int = 384,
+    patch_size: int = 16,
+):
+    generated_tokens = torch.zeros((parallel_size, image_token_num_per_image), dtype=torch.int).cuda()
+
+    inputs_embeds = prepare_classifier_free_guidance_input(inputs_embeds, vl_chat_processor, mmgpt, parallel_size)
+
+    for i in range(image_token_num_per_image):
+        outputs = mmgpt.language_model.model(inputs_embeds=inputs_embeds, use_cache=True, past_key_values=outputs.past_key_values if i != 0 else None)
+        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.unsqueeze(dim=1), next_token.unsqueeze(dim=1)], dim=1).view(-1)
+        img_embeds = mmgpt.prepare_gen_img_embeds(next_token)
+        inputs_embeds = img_embeds.unsqueeze(dim=1)
+
+    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)
+
+    visual_img = np.zeros((parallel_size, img_size, img_size, 3), dtype=np.uint8)
+    visual_img[:, :, :] = dec
+
+    generated_images = []
+    for i in range(parallel_size):
+        generated_images.append(PIL.Image.fromarray(visual_img[i]))
+
+    return generated_images
+
+@lru_cache(maxsize=1)
+def get_start_tag_embed(vl_gpt, vl_chat_processor):
+    with torch.no_grad():
+        return vl_gpt.language_model.get_input_embeddings()(
+            vl_chat_processor.tokenizer.encode(vl_chat_processor.image_start_tag, add_special_tokens=False, return_tensors="pt").to(vl_gpt.device)
+        )
+    
+def process_and_generate(vl_gpt, vl_chat_processor, input_image, prompt, num_images=4, cfg_weight=5):
+    start_tag_embed = get_start_tag_embed(vl_gpt, vl_chat_processor)
+
+    nl = '\n'
+    conversation = [
+        {
+            "role": "User",
+            "content": f"<image_placeholder>{nl + prompt if prompt else ''}",
+            "images": [input_image],
+        },
+        {"role": "Assistant", "content": ""},
+    ]
+
+    pil_images = load_pil_images(conversation)
+    prepare_inputs = vl_chat_processor(
+        conversations=conversation, images=pil_images, force_batchify=True
+    ).to(vl_gpt.device)
+
+    with torch.no_grad():
+        inputs_embeds = vl_gpt.prepare_inputs_embeds(**prepare_inputs)
+
+    inputs_embeds = torch.cat((inputs_embeds, start_tag_embed), dim=1)
+
+    generated_images = generate(
+        vl_gpt,
+        vl_chat_processor,
+        inputs_embeds,
+        parallel_size=num_images,
+        cfg_weight=cfg_weight
+    )
+
+    return generated_images

model_loader.py

@@ -0,0 +1,14 @@
+import torch
+from transformers import AutoModelForCausalLM
+from janus.models import MultiModalityCausalLM, VLChatProcessor
+
+def load_model_and_processor(model_path):
+    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.to(torch.bfloat16).cuda().eval()
+
+    return vl_gpt, vl_chat_processor

requirements.txt

@@ -0,0 +1,7 @@
+torch
+numpy
+Pillow
+gradio
+janus @ git+https://github.com/deepseek-ai/Janus
+transformers
+spaces