Trying out Janus-Pro on Google Colab

I have summarized my experience trying out "Janus-Pro" on "Google Colab".

1. Janus

1-1. Janus-Pro

"Janus-Pro" is an evolved version of its predecessor, "Janus". Specifically, "Janus-Pro" incorporates (1) optimized training strategies, (2) expanded training data, and (3) scaling to larger model sizes. With these improvements, "Janus-Pro" has achieved significant progress in both "multimodal understanding" and "Text-to-Image" instruction following capabilities, and has also improved the stability of "Text-to-Image" generation.

1-2. Janus

"Janus" is a new autoregressive framework that integrates multimodal understanding and generation. While it uses a single unified Transformer architecture for processing, it addresses the limitations of traditional approaches by separating visual encoding into distinct paths. This separation not only reduces conflicts in the roles of visual encoders in understanding and generation, but also improves the flexibility of the framework. "Janus" outperforms previous unified models and rivals or exceeds the performance of task-specific models. The simplicity, high flexibility, and effectiveness of "Janus" make it a strong candidate for the next generation of unified multimodal models.

1-3. JanusFlow

"JanusFlow" introduces a minimal architecture that integrates an autoregressive language model with rectified flow, a state-of-the-art method for generative modeling. The key finding shows that rectified flow can be easily learned within a large-scale language model framework without the need for complex architectural changes. Extensive experiments demonstrate that "JanusFlow" achieves performance equal to or better than specialized models in their respective domains, and significantly outperforms existing unified approaches across standard benchmarks. This research marks a step toward more efficient and versatile vision-language models.

2. Models

The following four "Janus" models are provided.

3. Image Generation

The procedure for executing image generation on "Google Colab" is as follows.

(1) Install packages.

# パッケージのインストール
!git clone https://github.com/deepseek-ai/Janus
%cd Janus
!pip install -e .

(2) Image generation.
This time, I will use "Janus-Pro-7B".

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


# specify the path to the model
model_path = "deepseek-ai/Janus-Pro-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
)
vl_gpt = vl_gpt.to(torch.bfloat16).cuda().eval()

conversation = [
    {
        "role": "<|User|>",
        "content": "A stunning princess from japan in red, white traditional clothing, black eyes, black hair of japanese anime style",
    },
    {"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


@torch.inference_mode()
def generate(
    mmgpt: MultiModalityCausalLM,
    vl_chat_processor: VLChatProcessor,
    prompt: str,
    temperature: float = 1,
    parallel_size: int = 16,
    cfg_weight: float = 5,
    image_token_num_per_image: int = 576,
    img_size: int = 384,
    patch_size: int = 16,
):
    input_ids = vl_chat_processor.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] = vl_chat_processor.pad_id

    inputs_embeds = mmgpt.language_model.get_input_embeddings()(tokens)

    generated_tokens = torch.zeros((parallel_size, image_token_num_per_image), dtype=torch.int).cuda()

    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

    os.makedirs('generated_samples', exist_ok=True)
    for i in range(parallel_size):
        save_path = os.path.join('generated_samples', "img_{}.jpg".format(i))
        PIL.Image.fromarray(visual_img[i]).save(save_path)


generate(
    vl_gpt,
    vl_chat_processor,
    prompt,
)

16 images have been generated under "Janus/generate_samples".

The memory consumption is as follows.