first commit

src/utils/generation_utils.py

+# -*- coding: utf-8 -*-
+# Copyright 2025 BAAI. and/or its affiliates.
+# SPDX-License-Identifier: Apache-2.0
+
+import re
+from PIL import Image
+
+import numpy as np
+import torch
+from transformers import GenerationConfig
+from transformers.generation import LogitsProcessorList
+
+from .logits_processor import UnbatchedClassifierFreeGuidanceLogitsForVisualTokenWithDifferentialTopKProcessor
+
+
+@torch.no_grad()
+def generate(
+    cfg,
+    model,
+    tokenizer,
+    input_ids,
+    unconditional_ids,
+    full_unconditional_ids=None,
+    force_same_image_size=True,
+):
+    if cfg.streaming:
+        raise NotImplementedError("Not supported streaming generation")
+        # yield from streaming_generate(cfg, model, tokenizer, input_ids, unconditional_ids, full_unconditional_ids)
+    else:
+        yield non_streaming_generate(cfg, model, tokenizer, input_ids, unconditional_ids, full_unconditional_ids, force_same_image_size)
+
+
+def streaming_generate(
+    cfg,
+    model,
+    tokenizer,
+    input_ids,
+    unconditional_ids,
+    full_unconditional_ids=None,
+):
+    pass
+
+
+def non_streaming_generate(
+    cfg,
+    model,
+    tokenizer,
+    input_ids,
+    unconditional_ids,
+    full_unconditional_ids=None,
+    force_same_image_size=True,
+):
+    input_ids_len = input_ids.shape[1]
+
+    logits_processor = LogitsProcessorList()
+    logits_processor.append(
+        build_logits_processor(
+            cfg,
+            unconditional_ids,
+            model,
+            tokenizer,
+            full_unconditional_ids,
+            force_same_image_size=force_same_image_size,
+        )
+    )
+
+    generation_config = GenerationConfig(
+        **cfg.sampling_params,
+        pad_token_id=cfg.special_token_ids["PAD"],
+        eos_token_id=cfg.special_token_ids["EOS"],
+    )
+
+    token_ids = model.generate(
+        input_ids,
+        generation_config,
+        logits_processor=logits_processor,
+    )
+
+    gen_token_ids = token_ids[:, input_ids_len:]
+    return gen_token_ids[0].detach().cpu().numpy()
+
+
+def build_logits_processor(
+    cfg,
+    unconditional_ids,
+    model,
+    tokenizer,
+    full_unconditional_ids=None,
+    force_same_image_size=True,
+):
+    logits_processor = UnbatchedClassifierFreeGuidanceLogitsForVisualTokenWithDifferentialTopKProcessor(
+        guidance_scale=cfg.classifier_free_guidance,
+        unconditional_ids=unconditional_ids,
+        full_unconditional_ids=full_unconditional_ids,
+        model=model,
+        tokenizer=tokenizer,
+        unconditional_type=cfg.unconditional_type,
+        target_height=getattr(cfg, "target_height", None),
+        target_width=getattr(cfg, "target_width", None),
+        image_cfg_scale=getattr(cfg, "image_cfg_scale", 1.0),
+        use_differential_sampling=cfg.sampling_params["use_differential_sampling"],
+        text_top_k=cfg.sampling_params["text_top_k"],
+        text_top_p=cfg.sampling_params["text_top_p"],
+        text_temperature=cfg.sampling_params["text_temperature"],
+        image_top_k=cfg.sampling_params["image_top_k"],
+        image_top_p=cfg.sampling_params["image_top_p"],
+        image_temperature=cfg.sampling_params["image_temperature"],
+        force_same_image_size=force_same_image_size,
+    )
+
+    return logits_processor
+
+
+@torch.no_grad()
+def multimodal_decode(
+    outputs,
+    tokenizer,
+    vision_tokenizer,
+):
+    outputs = outputs.replace("<|extra_101|>", "").replace("<|extra_204|>", "")
+    pattern = re.compile(
+        rf"({re.escape(tokenizer.bog_token)}.*?{re.escape(tokenizer.eog_token)}|"
+        rf"{re.escape(tokenizer.boc_token)}.*?{re.escape(tokenizer.eoc_token)}|"
+        rf"{re.escape(tokenizer.boi_token)}.*?{re.escape(tokenizer.eoi_token)})",
+        re.DOTALL
+    )
+
+    multimodal_output = []
+    chunks = re.split(pattern, outputs)
+    for c in chunks:
+        if len(c) == 0:
+            continue
+
+        if tokenizer.boi_token in c and tokenizer.eoi_token in c:
+            image = decode_image(c, tokenizer, vision_tokenizer)
+            if image is not None:
+                multimodal_output.append(("image", image))
+        elif tokenizer.bog_token in c and tokenizer.eog_token in c:
+            multimodal_output.append(
+                ("global_cot", c.replace(tokenizer.bog_token, "").replace(tokenizer.eog_token, ""))
+            )
+        elif tokenizer.boc_token in c and tokenizer.eoc_token in c:
+            multimodal_output.append(
+                ("image_cot", c.replace(tokenizer.boc_token, "").replace(tokenizer.eoc_token, ""))
+            )
+        # exclude incomplete image
+        elif tokenizer.boi_token not in c and len(c.strip()) > 0:
+            multimodal_output.append(("text", c))
+
+    return multimodal_output
+
+
+def decode_image(image_string, tokenizer, vision_tokenizer):
+    image = []
+    image_rows = re.split(re.escape(tokenizer.eol_token), image_string)
+    for r in image_rows:
+        token_ids = re.findall(r"<\|visual token (\d+)\|>", r)
+        if len(token_ids) > 0:
+            row_token = [int(m) for m in token_ids]
+            image.append(row_token)
+
+    try:
+        image = torch.tensor(image, dtype=torch.long, device=next(iter(vision_tokenizer.parameters())).device)
+        h, w = image.shape
+        image = vision_tokenizer.decode_code(image[None], shape=(1, h, w, 256)).float()
+        image = image[0].permute(1, 2, 0)
+        image = Image.fromarray(((image + 1.0) * 127.5).clamp(0, 255).detach().cpu().numpy().astype(np.uint8))
+        return image
+    except Exception as ex:
+        print(f"decode image failed {ex}")
+        return None

src/utils/input_utils.py

+# -*- coding: utf-8 -*-
+# Copyright 2025 BAAI. and/or its affiliates.
+# SPDX-License-Identifier: Apache-2.0
+
+from PIL import Image
+import torch
+import numpy as np
+
+
+def smart_resize(image: Image.Image, area: int = 512 * 512, ds_factor: int = 16):
+    width, height = image.size
+    aspect_ratio = width / height
+    new_height = int((area / aspect_ratio) ** 0.5)
+    new_width = int(new_height * aspect_ratio)
+    # Round to nearest multiple of divisible_by
+    new_height = ((new_height + ds_factor//2) // ds_factor) * ds_factor
+    new_width = ((new_width + ds_factor//2) // ds_factor) * ds_factor
+    return image.resize((new_width, new_height), Image.BICUBIC)
+
+
+def format_image_string(tokenizer, image_tokens):
+    image_string = ""
+    h, w = image_tokens.shape
+    for _h in range(h):
+        row_string = ""
+        for _w in range(w):
+            row_string += "<|visual token {token_id:0>6d}|>".format(token_id=image_tokens[_h, _w])
+
+        if _h < h - 1:
+            row_string += tokenizer.eol_token
+        image_string += row_string
+
+    return "{image_start}{token_height}*{token_width}{image_token}{token_str}{image_end}".format(
+        image_start=tokenizer.boi_token,
+        token_height=h,
+        token_width=w,
+        image_token=tokenizer.img_token,
+        token_str=image_string,
+        image_end=tokenizer.eoi_token,
+    )
+
+
+@torch.no_grad()
+def build_image(image, cfg, tokenizer, vq_model):
+    image = smart_resize(image, cfg.image_area)
+    w, h = image.size
+    device = next(vq_model.parameters()).device
+    dtype = next(vq_model.parameters()).dtype
+    image = torch.tensor((np.array(image) / 127.5 - 1.0)).to(device, dtype).permute(2, 0, 1)
+    _, _, token = vq_model.encode(image[None])
+    token = token[-1].view(h // 16, w // 16)
+    return format_image_string(tokenizer, token)

src/utils/logging_utils.py

+# -*- coding:utf-8 -*-
+# Copyright 2025 BAAI. and/or its affiliates.
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+from datetime import datetime
+import os.path as osp
+import builtins
+old_print = builtins.print
+
+
+def setup_print_file(file):
+    def print(*args, **kwargs):
+        msg = " ".join(map(str, args))
+        with open(file, "a") as f:
+            f.write(msg + "\n")
+        old_print(msg)
+
+    builtins.print = print
+
+
+def setup_logger(log_dir="./", log_name="log"):
+    logfile = osp.join(
+        log_dir,
+        f'{log_name}_{datetime.now().strftime("%Y%m%d_%H%M%S")}.log',
+    )
+    os.makedirs(osp.dirname(logfile), exist_ok=True)
+    setup_print_file(logfile)

src/utils/logits_processor.py

+# Copyright 2025 BAAI. and/or its affiliates.
+# SPDX-License-Identifier: Apache-2.0
+
+import math
+import os
+from typing import Callable, Dict, Iterable, List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+from torch.nn import functional as F
+from transformers.generation import LogitsProcessor
+
+BOS = 151849
+EOS = 151850
+IMG = 151851
+BOI = 151852
+EOI = 151853
+EOL = 151846
+EOF = 151847
+BOV = 151854
+
+
+class UnbatchedClassifierFreeGuidanceLogitsForVisualTokenProcessor(LogitsProcessor):
+
+    def __init__(
+        self,
+        guidance_scale: float,
+        model,
+        tokenizer,
+        unconditional_ids: torch.LongTensor,
+        unconditional_attention_mask: Optional[torch.LongTensor] = None,
+        full_unconditional_ids: Optional[torch.LongTensor] = None,
+        full_unconditional_attention_mask: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = True,
+        allowed_tokens_control: bool = True,
+        force_same_image_size: bool = True,
+        unconditional_type: str = "no_text", # options: no_text, no_prev_text, no_prev_modal, no_text_img_cfg, etc.
+        target_height: Optional[int] = None,  # added parameter is used to specify the target height
+        target_width: Optional[int] = None,   # added parameter is used to specify the target width
+        image_cfg_scale: float = 1.0,
+    ):
+        self.guidance_scale = guidance_scale
+        self.model = model
+        self.tokenizer = tokenizer
+
+        self.unconditional_context = {
+            "input_ids": unconditional_ids,
+            "attention_mask": (
+                unconditional_attention_mask
+                if unconditional_attention_mask is not None
+                else torch.ones_like(unconditional_ids, dtype=torch.long)
+            ),
+            "default_input_ids": unconditional_ids,
+            "default_attention_mask": (
+                unconditional_attention_mask
+                if unconditional_attention_mask is not None
+                else torch.ones_like(unconditional_ids, dtype=torch.long)
+            ),
+            "first_pass": True,
+            "past_key_values": None,
+        }
+
+        self.image_cfg_scale = image_cfg_scale
+        # heuristic combination example: fuc + a*(uc - fuc) + b*(logit - un)
+        self.full_unconditional_context = {
+            "input_ids": full_unconditional_ids if full_unconditional_ids is not None else unconditional_ids,
+            "attention_mask": (
+                unconditional_attention_mask
+                if unconditional_attention_mask is not None
+                else torch.ones_like(unconditional_ids, dtype=torch.long)
+            ),
+            "default_input_ids": full_unconditional_ids if full_unconditional_ids is not None else unconditional_ids,
+            "default_attention_mask": (
+                full_unconditional_attention_mask
+                if full_unconditional_attention_mask is not None
+                else (
+                    unconditional_attention_mask
+                    if unconditional_attention_mask is not None
+                    else
+                    torch.ones_like(unconditional_ids, dtype=torch.long)
+                )
+            ),
+            "first_pass": True,
+            "past_key_values": None,
+        }
+
+        self.use_cache = use_cache
+
+        self.first_in_image = True
+        self.in_image = False
+        self.in_visual = False
+        self.image_nums = 0
+
+        self.allowed_tokens_control = allowed_tokens_control
+
+        self.height = target_height 
+        self.width = target_width
+        if self.height is not None and self.width is not None:
+            print(f"[INFO] User defined: height: {self.height}, width: {self.width}")
+        else:
+            print(f"[INFO] Auto height, width")
+        self.hw_tokens = None
+        self.force_same_image_size = force_same_image_size
+        self.unconditional_type = unconditional_type
+
+        self.text_segment = 0
+
+        if IMG in unconditional_ids[0]:
+            self.parse_hw(unconditional_ids)
+            self.first_in_image = False
+
+    # used only for visual tokens with classifier-free guidance
+    def get_unconditional_logits(self, input_ids):
+        # update cache
+        input_ids, attention_mask = self.update_unconditional_context(input_ids, target="unconditional")
+        if self.use_cache and not self.unconditional_context["first_pass"]:
+            input_ids = input_ids[:, -1:]
+        else:
+            self.unconditional_context["first_pass"] = False
+
+        out = self.model(
+            input_ids,
+            attention_mask=attention_mask,
+            use_cache=self.use_cache,
+            past_key_values=self.unconditional_context["past_key_values"],
+        )
+
+        if self.use_cache:
+            self.unconditional_context["past_key_values"] = out.get("past_key_values", None)
+
+        if self.unconditional_type == "no_text_img_cfg":
+            input_ids, attention_mask = self.update_unconditional_context(input_ids, target="full_unconditional")
+
+            if self.use_cache and not self.full_unconditional_context["first_pass"]:
+                input_ids = input_ids[:, -1:]
+            else:
+                self.full_unconditional_context["first_pass"] = False
+
+            func_out = self.model(
+                input_ids,
+                attention_mask=attention_mask,
+                use_cache=self.use_cache,
+                past_key_values=self.full_unconditional_context["past_key_values"],
+            )
+
+            if self.use_cache:
+                self.full_unconditional_context["past_key_values"] = out.get("past_key_values", None)
+
+            return out.logits[:, -1], func_out.logits[:, -1]
+
+        return out.logits[:, -1]
+
+    def __call__(self, input_ids, scores):
+        # IMAGE MODE
+        if input_ids[0][-1] == BOI:
+            self.set_unconditional_context(input_ids)
+        # TEXT MODE
+        if input_ids[0][-1] == EOI:
+            self.exit_image(input_ids)
+
+        if input_ids[0][-1] != BOI and input_ids[0][-2] == EOI:
+            self.text_segment += 1
+
+        # all image tokens between BOI and EOI are handled by in_image_logits_processor
+        N = input_ids.shape[1]
+        if self.in_image:
+            scores = self.in_image_logits_processor(input_ids, scores)
+        else:
+            scores = self.in_text_logits_processor(input_ids, scores)
+
+        return scores
+
+    def in_image_logits_processor(self, input_ids, scores):
+        # all visual-related logits must call get_unconditional_logits to keep unconditional cache updated
+        unc_scores = self.get_unconditional_logits(input_ids)
+
+        if self.in_visual:
+            # generating visual tokens
+            img_idx = self.find_last_token_index(input_ids[0], IMG)
+            vis_idx = input_ids.shape[1] - img_idx
+            # eoi
+            if vis_idx == self.height * (self.width + 1):
+                mask = torch.full_like(scores, -math.inf)
+                mask[:, EOI] = 0
+                scores = scores + mask
+                return scores
+            # eol
+            elif vis_idx % (self.width + 1) == 0:
+                mask = torch.full_like(scores, -math.inf)
+                mask[:, EOL] = 0
+                scores = scores + mask
+                return scores
+            # visual token
+            else:
+                scores = self.apply_cfg(scores, unc_scores, self.guidance_scale, self.image_cfg_scale)
+                mask = torch.full_like(scores, -math.inf)
+                mask[:, BOV:] = 0
+                scores = scores + mask
+                return scores
+        else:
+            # height and width have been generated; confirm h and w
+            if input_ids[0][-1] == IMG:
+                self.in_visual = True
+                self.image_nums += 1
+                print(f"[INFO] Generating image#{self.image_nums}...")
+                if self.first_in_image or not self.force_same_image_size:
+                    if self.height is None or self.width is None:  # Only when the size is not specified will the analysis be conducted.
+                        self.parse_hw(input_ids)
+                    self.first_in_image = False
+                
+                # the first visual token
+                scores = self.apply_cfg(scores, unc_scores, self.guidance_scale, self.image_cfg_scale)
+
+                # mask non visual tokens
+                mask = torch.full_like(scores, -math.inf)
+                mask[:, BOV:] = 0
+                scores = scores + mask
+                return scores
+            # not yet at IMG token: we are generating h x w tokens
+            else:
+                if self.height is not None and self.width is not None:
+                    # If the size is specified, the corresponding h and w tokens will be forcibly generated.
+                    boi_idx = self.find_last_token_index(input_ids[0], BOI)
+                    hw_idx = input_ids.shape[1] - boi_idx
+                    hw_tokens = self.tokenizer.encode(f"{str(self.height)}*{str(self.width)}", add_special_tokens=False)
+
+                    if hw_idx <= len(hw_tokens):  # height part
+                        allowed_tokens = [hw_tokens[hw_idx-1]]
+                    else:  # IMG token
+                        allowed_tokens = [IMG]
+                        
+                    mask = torch.full_like(scores, -math.inf)
+                    for token in allowed_tokens:
+                        mask[:, token] = 0
+                    scores = scores + mask
+                    return scores
+                else:
+                    if self.first_in_image or not self.force_same_image_size:
+                        # force output format for h x w sequence
+                        boi_idx = self.find_last_token_index(input_ids[0], BOI)
+                        hw_idx = input_ids.shape[1] - boi_idx
+                        mask = torch.full_like(scores, -math.inf)
+                        if hw_idx > 5:
+                            mask[:, IMG] = 0
+                        elif hw_idx in (1, 4):
+                            # example mapping: '1' -> token 16, '9' -> token 24
+                            mask[:, 16:25] = 0
+                        elif hw_idx == 3:
+                            # '*' -> token 9
+                            mask[:, 9] = 0
+                        elif hw_idx in (2, 5):
+                            # '0' -> token 15, '9' -> token 24
+                            mask[:, 15:25] = 0
+                        scores = scores + mask
+                        return scores
+                    else:
+                        # force same image size: must reproduce previous h and w tokens
+                        boi_idx = self.find_last_token_index(input_ids[0], BOI)
+                        hw_idx = input_ids.shape[1] - boi_idx
+                        mask = torch.full_like(scores, -math.inf)
+                        if hw_idx > len(self.hw_tokens):
+                            mask[:, IMG] = 0
+                        else:
+                            mask[:, self.hw_tokens[hw_idx - 1]] = 0
+                        scores = scores + mask
+                        return scores
+
+    def in_text_logits_processor(self, input_ids, scores):
+        # for text tokens, do not apply CFG here
+        mask = torch.full_like(scores, -math.inf)
+        mask[:, :BOV] = 0
+        scores = scores + mask
+        return scores
+
+    def parse_hw(self, input_ids):
+        if self.height is not None and self.width is not None and self.force_same_image_size:
+            return
+
+        # Find indices of last BOI and IMG tokens in the sequence
+        seq = input_ids[0]
+        img_indices = (seq == IMG).nonzero().flatten()
+        boi_indices = (seq == BOI).nonzero().flatten()
+
+        # Get the last occurrence of each token
+        last_img_pos = img_indices[-1].item()
+        last_boi_pos = boi_indices[-1].item()
+
+        # Get tokens between last BOI and IMG
+        self.hw_tokens = seq[last_boi_pos+1:last_img_pos]
+
+        # Decode tokens to string and parse dimensions
+        hw_str = self.tokenizer.decode(self.hw_tokens)
+        h, w = hw_str.split('*')
+        self.height = int(h)
+        self.width = int(w)
+        print(f"[INFO] hw_str: {hw_str}, H_px: {self.height * 16}, W_px: {self.width * 16}")
+
+    def find_last_token_index(self, seq, token_id):
+        # seq shape: (N,)
+        token_indices = (seq == token_id).nonzero().flatten()
+        if len(token_indices) == 0:
+            return -1
+        return token_indices[-1].item()
+
+    def find_first_token_index(self, seq, token_id):
+        token_indices = (seq == token_id).nonzero().flatten()
+        if len(token_indices) == 0:
+            return -1
+        return token_indices[0].item()
+
+    def apply_cfg(self, scores, unc_scores, cfg_scale, img_cfg_scale):
+        match self.unconditional_type:
+            case "no_text":
+                # applying no_text CFG
+                scores = F.log_softmax(scores, dim=-1)
+                unc_scores = F.log_softmax(unc_scores, dim=-1)
+                # c = unc + cfg * (c - unc)
+                scores = cfg_scale * (scores - unc_scores) + unc_scores
+                return scores
+            case "no_text_img_cfg":
+                # applying no_text and no_any CFG
+                no_text_scores, no_any_scores = unc_scores
+                scores = F.log_softmax(scores, dim=-1)
+                no_text_scores = F.log_softmax(no_text_scores, dim=-1)
+                no_any_scores = F.log_softmax(no_any_scores, dim=-1)
+                # c = unc + cfg_1 * (unc_im - unc) + cfg_2 * (c - unc_im)
+                scores = no_any_scores + img_cfg_scale * (no_text_scores - no_any_scores) + cfg_scale * (scores - no_text_scores)
+                return scores
+            case _:
+                raise NotImplementedError(f"Unsupported unconditional type {self.unconditional_type}")
+
+    def set_unconditional_context(self, input_ids):
+        self.in_image = True
+
+        if self.unconditional_type == "no_text" or self.unconditional_type == "uncondition_tokens":
+            # no special handling
+            return
+        elif self.unconditional_type == "no_text_img_cfg":
+            self.full_unconditional_context["input_ids"] = self.full_unconditional_context["default_input_ids"]
+            self.full_unconditional_context["attention_mask"] = self.full_unconditional_context["default_attention_mask"]
+            self.full_unconditional_context["first_pass"] = True
+            self.full_unconditional_context["past_key_values"] = None
+        else:
+            raise ValueError(f"Unconditional type {self.unconditional_type} not supported")
+
+    def update_unconditional_context(self, input_ids, target="unconditional"):
+        if target == "unconditional":
+            self.unconditional_context["input_ids"] = torch.cat(
+                [
+                    self.unconditional_context["input_ids"],
+                    input_ids[:, -1:],
+                ],
+                dim=1,
+            )
+            self.unconditional_context["attention_mask"] = torch.cat(
+                [
+                    self.unconditional_context["attention_mask"],
+                    torch.ones_like(input_ids[:, -1:], dtype=torch.long),
+                ],
+                dim=1,
+            )
+            return self.unconditional_context["input_ids"], self.unconditional_context["attention_mask"]
+        elif target == "full_unconditional":
+            self.full_unconditional_context["input_ids"] = torch.cat(
+                [
+                    self.full_unconditional_context["input_ids"],
+                    input_ids[:, -1:],
+                ],
+                dim=1,
+            )
+            self.full_unconditional_context["attention_mask"] = torch.cat(
+                [
+                    self.full_unconditional_context["attention_mask"],
+                    torch.ones_like(input_ids[:, -1:], dtype=torch.long),
+                ],
+                dim=1,
+            )
+            return self.full_unconditional_context["input_ids"], self.full_unconditional_context["attention_mask"]
+
+
+    def exit_image(self, input_ids):
+        self.in_image = False
+        self.in_visual = False
+        # update eoi to unconditional cache
+        self.get_unconditional_logits(input_ids)
+
+
+class UnbatchedClassifierFreeGuidanceLogitsForVisualTokenWithDifferentialTopKProcessor(UnbatchedClassifierFreeGuidanceLogitsForVisualTokenProcessor):
+    """
+    Extend the original CFG processor and add the differentiated TopK function.
+    Use different topk, top_p, and temperature parameters when generating text tokens and image tokens.
+    """
+
+    def __init__(
+        self,
+        guidance_scale: float,
+        model,
+        tokenizer,
+        unconditional_ids: Optional[torch.LongTensor],
+        unconditional_attention_mask: Optional[torch.LongTensor] = None,
+        full_unconditional_ids: Optional[torch.LongTensor] = None,
+        full_unconditional_attention_mask: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = True,
+        allowed_tokens_control: bool = True,
+        force_same_image_size: bool = True,
+        unconditional_type: str = "no_text",
+        target_height: Optional[int] = None,
+        target_width: Optional[int] = None,
+        image_cfg_scale: float = 1.0,
+        # added differentiated topk parameter
+        use_differential_sampling: bool = False,
+        text_top_k: int = 1024,
+        image_top_k: int = 10240,
+        text_top_p: float = 0.9,
+        image_top_p: float = 1.0,
+        text_temperature: float = 1.0,
+        image_temperature: float = 1.0,
+        **kwargs,
+    ):
+
+        super().__init__(
+            guidance_scale=guidance_scale,
+            model=model,
+            tokenizer=tokenizer,
+            unconditional_ids=unconditional_ids,
+            unconditional_attention_mask=unconditional_attention_mask,
+            full_unconditional_ids=full_unconditional_ids,
+            full_unconditional_attention_mask=full_unconditional_attention_mask,
+            use_cache=use_cache,
+            allowed_tokens_control=allowed_tokens_control,
+            force_same_image_size=force_same_image_size,
+            unconditional_type=unconditional_type,
+            target_height=target_height,
+            target_width=target_width,
+            **kwargs,
+        )
+
+        # differential sampling parameters
+        self.use_differential_sampling = use_differential_sampling
+        self.text_top_k = text_top_k
+        self.image_top_k = image_top_k
+        self.text_top_p = text_top_p
+        self.image_top_p = image_top_p
+        self.text_temperature = text_temperature
+        self.image_temperature = image_temperature
+
+    def apply_differential_topk(self, scores, is_image_generation=False):
+        """Apply mode-specific top_k/top_p/temperature depending on image/text generation."""
+        if not self.use_differential_sampling:
+            return scores
+
+        # choose params based on current mode
+        if is_image_generation:
+            current_top_k = self.image_top_k
+            current_top_p = self.image_top_p
+            current_temperature = self.image_temperature
+            mode = "IMAGE"
+        else:
+            current_top_k = self.text_top_k
+            current_top_p = self.text_top_p
+            current_temperature = self.text_temperature
+            mode = "TEXT"
+
+        # temperature scaling
+        if current_temperature != 1.0:
+            scores = scores / current_temperature
+
+        # Top-K filtering
+        if current_top_k > 0 and current_top_k < scores.size(-1):
+            top_k = min(current_top_k, scores.size(-1))
+            # mask everything below kth logit
+            indices_to_remove = scores < torch.topk(scores, top_k)[0][..., -1, None]
+            scores[indices_to_remove] = float('-inf')
+
+        # Top-P (nucleus) filtering
+        if current_top_p < 1.0:
+            sorted_logits, sorted_indices = torch.sort(scores, descending=True)
+            cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
+
+            # find positions exceeding cumulative probability threshold
+            sorted_indices_to_remove = cumulative_probs > current_top_p
+            # keep first token above threshold
+            sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+            sorted_indices_to_remove[..., 0] = 0
+
+            # map mask back to original indices
+            indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
+            scores[indices_to_remove] = float('-inf')
+
+        return scores
+
+    def in_image_logits_processor(self, input_ids, scores):
+        """Override: add differential top-k for image tokens."""
+        # apply base processing first
+        scores = super().in_image_logits_processor(input_ids, scores)
+
+        # then apply differential top-k for image tokens
+        scores = self.apply_differential_topk(scores, is_image_generation=True)
+
+        return scores
+
+    def in_text_logits_processor(self, input_ids, scores):
+        """Override: add differential top-k for text tokens."""
+        # apply base processing first
+        scores = super().in_text_logits_processor(input_ids, scores)
+
+        # then apply differential top-k for text tokens
+        scores = self.apply_differential_topk(scores, is_image_generation=False)
+
+        return scores
\ No newline at end of file

src/utils/model_utils.py

+# -*- coding: utf-8 -*-
+# Copyright 2025 BAAI. and/or its affiliates.
+# SPDX-License-Identifier: Apache-2.0
+
+import os.path as osp
+import torch
+from transformers import AutoTokenizer
+
+from ..emu3p5 import Emu3ForCausalLM, Emu3Config
+from ..vision_tokenizer import build_vision_tokenizer
+
+def build_emu3p5(
+    model_path,
+    tokenizer_path,
+    vq_path,
+    vq_type="ibq",
+    model_device="auto",
+    vq_device="cuda:0",
+    **kwargs,
+):
+    if isinstance(model_device, int):
+        device_map = f"cuda:{model_device}"
+    else:
+        device_map = model_device
+
+    print(device_map)
+
+    # MLLM
+    model_config = Emu3Config.from_pretrained(
+        model_path,
+        trust_remote_code=True,
+    )
+    model = Emu3ForCausalLM.from_pretrained(
+        model_path,
+        config=model_config,
+        torch_dtype=torch.bfloat16,
+        device_map=device_map,
+        attn_implementation="flash_attention_2",
+        # attn_implementation="eager", # if you cann't install flash_attention
+    )
+    model.eval()
+    
+    # text tokenizer
+    tokenizer = AutoTokenizer.from_pretrained(
+        tokenizer_path,
+        special_tokens_file=osp.join(tokenizer_path, "emu3_vision_tokens.txt"),
+        trust_remote_code=True,
+    )
+    tokenizer.bos_token = "<|extra_203|>"
+    tokenizer.eos_token = "<|extra_204|>"
+    tokenizer.pad_token = "<|endoftext|>"
+    tokenizer.eol_token = "<|extra_200|>"
+    tokenizer.eof_token = "<|extra_201|>"
+    tokenizer.tms_token = "<|extra_202|>"
+    tokenizer.img_token = "<|image token|>"
+    tokenizer.boi_token = "<|image start|>"
+    tokenizer.eoi_token = "<|image end|>"
+    tokenizer.bss_token = "<|extra_100|>"
+    tokenizer.ess_token = "<|extra_101|>"
+    tokenizer.bog_token = "<|extra_60|>"
+    tokenizer.eog_token = "<|extra_61|>"
+    tokenizer.boc_token = "<|extra_50|>"
+    tokenizer.eoc_token = "<|extra_51|>"
+
+    # vq tokenizer
+    vq_model = build_vision_tokenizer(vq_type, vq_path, device=vq_device, **kwargs)
+
+    return model, tokenizer, vq_model
+

src/utils/painting_utils.py

+# -*- coding: utf-8 -*-
+# Copyright 2025 BAAI. and/or its affiliates.
+# SPDX-License-Identifier: Apache-2.0
+
+import io
+import re
+import os.path as osp
+
+from PIL import Image
+
+from src.proto import emu_pb as story_pb
+
+class ProtoWriter:
+    
+    def __init__(self):
+        self.story = story_pb.Story()
+        self.image_tensor = None
+
+    def clear(self):
+        self.story = story_pb.Story()
+        self.image_tensor = None
+
+    def extend(self, multimodal_output):
+        for t, c in multimodal_output:
+            match t:
+                case "question":
+                    self.story.question = c
+                case "global_cot":
+                    self.story.summary = c
+                case "image_cot":
+                    image = story_pb.ImageMeta()
+                    image.chain_of_thought = c
+                    self._put_last_image(image)
+                case "text":
+                    self._put_last_clip(self._build_clip(c))
+                case "image":
+                    image = self._get_last_image()
+                    image.image.CopyFrom(self._build_image(c))
+                    self._put_last_image(image)
+                case "reference_image":
+                    image = story_pb.ImageMeta()
+                    image.image.CopyFrom(self._build_image(c))
+                    self.story.reference_images.append(image)
+                case _:
+                    raise NotImplementedError(f"Unsupported data type {t}")
+
+    def save(self, path):
+        self._check_last_image()
+        with open(path, 'wb') as f:
+            f.write(self.story.SerializeToString())
+
+
+    def _build_clip(self, text_content=""):
+        clip = story_pb.Clip()
+        clip.clip_id = f"clip_{len(self.story.clips):04d}"
+        segment = story_pb.Segment()
+        segment.asr = text_content
+
+        clip.segments.append(segment)
+        return clip
+
+    def _build_image(self, image):
+        im = story_pb.Image()
+        im.width, im.height = image.size
+        im.format = story_pb.ImageFormat.PNG
+
+        img_byte_arr = io.BytesIO()
+        image.save(img_byte_arr, format="PNG")
+        im.image_data = img_byte_arr.getvalue()
+
+        return im
+
+    def _get_last_image(self):
+        if not self.story.clips:
+            self._put_last_clip(self._build_clip())
+
+        if self.story.clips[-1].segments[0].images and not self.story.clips[-1].segments[0].images[-1].image.image_data:
+            image = self.story.clips[-1].segments[0].images[-1]
+            del self.story.clips[-1].segments[0].images[-1]
+        else:
+            image = story_pb.ImageMeta()
+
+        return image
+
+    def _put_last_image(self, image):
+        if not self.story.clips:
+            self._put_last_clip(self._build_clip())
+
+        self.story.clips[-1].segments[0].images.append(image)
+
+    def _put_last_clip(self, clip):
+        self.story.clips.append(clip)
+
+    def _check_last_image(self):
+        image = self._get_last_image()
+        if image.image.image_data:
+            self._put_last_image(image)

src/utils/video_utils.py

+# -*- coding: utf-8 -*-
+# Copyright 2025 BAAI. and/or its affiliates.
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+import os.path as osp
+import numpy as np
+import torch
+from PIL import Image, ImageDraw, ImageFont
+import imageio
+
+
+def wrap_text(draw, text, font, max_width):
+    lines = []
+    current_line = ""
+    
+    i = 0
+    while i < len(text):
+        char = text[i]
+        test_line = current_line + char
+        
+        bbox = draw.textbbox((0, 0), test_line, font=font)
+        text_width = bbox[2] - bbox[0]
+        
+        if text_width <= max_width:
+            current_line = test_line
+            i += 1
+        else:
+            if current_line:
+                lines.append(current_line)
+                current_line = ""
+            else:
+                current_line = char
+                i += 1
+    
+    if current_line:
+        lines.append(current_line)
+    
+    return lines
+
+
+def plot_string(string, font_path="src/proto/assets/cangerjinkai.ttf", font_size=80, image_size=(512, 512), bg_color="white", text_color="black"):
+    img = Image.new("RGB", image_size, color=bg_color)
+    draw = ImageDraw.Draw(img)
+
+    margin = 100
+    max_width = max(image_size[0] - 2 * margin, 1)
+    max_height = max(image_size[1] - 2 * margin, 1)
+
+    def load_font(size):
+        if font_path:
+            try:
+                return ImageFont.truetype(font_path, size)
+            except Exception:
+                print(f"Failed to load font from {font_path}")
+        return ImageFont.load_default()
+
+    font = load_font(font_size)
+    lines = wrap_text(draw, string, font, max_width)
+    line_height = draw.textbbox((0, 0), "Ay", font=font)[3]
+    total_text_height = line_height * max(len(lines), 1)
+
+    if total_text_height > max_height:
+        for size in range(font_size - 2, 9, -2):
+            font = load_font(size)
+            lines = wrap_text(draw, string, font, max_width)
+            line_height = draw.textbbox((0, 0), "Ay", font=font)[3]
+            total_text_height = line_height * max(len(lines), 1)
+            if total_text_height <= max_height:
+                break
+        else:
+            font = ImageFont.load_default()
+            lines = wrap_text(draw, string, font, max_width)
+            line_height = draw.textbbox((0, 0), "Ay", font=font)[3]
+            total_text_height = line_height * max(len(lines), 1)
+
+    y_offset = max(margin, (image_size[1] - total_text_height) // 2)
+
+    for line in lines:
+        bbox = draw.textbbox((0, 0), line, font=font)
+        text_width = bbox[2] - bbox[0]
+        x_offset = max(margin, (image_size[0] - text_width) // 2)
+        draw.text((x_offset, y_offset), line, fill=text_color, font=font)
+        y_offset += line_height
+
+    return np.array(img)
+
+
+def save_image_list_to_video(images, path, fps=1, quality='high'):
+    os.makedirs(osp.dirname(path), exist_ok=True)
+    
+    if '.mp4' not in path and len(images) == 1:
+        img = images[0]
+        if isinstance(img, torch.Tensor):
+            img = img.detach().cpu().numpy().astype(np.uint8)
+        elif isinstance(img, Image.Image):
+            img = np.array(img).astype(np.uint8)
+        else:
+            img = img.astype(np.uint8)
+        Image.fromarray(img).save(path, quality=100)
+        return
+    
+    func = lambda x: (
+        x.detach().cpu().numpy().astype(np.uint8)
+        if isinstance(x, torch.Tensor)
+        else x.astype(np.uint8)
+    )
+    images = list(map(func, images))
+    
+    if quality == 'high':
+        try:
+            writer = imageio.get_writer(
+                path,
+                fps=fps,
+                codec='libx264',
+                ffmpeg_params=[
+                    '-crf', '18',
+                    '-preset', 'slow',
+                    '-pix_fmt', 'yuv420p',
+                ]
+            )
+            for image in images:
+                writer.append_data(image)
+            writer.close()
+        except (TypeError, AttributeError):
+            try:
+                writer = imageio.get_writer(path, fps=fps, codec='libx264', macro_block_size=None)
+                for image in images:
+                    writer.append_data(image)
+                writer.close()
+            except Exception:
+                with imageio.get_writer(path, fps=fps, mode='I') as writer:
+                    for image in images:
+                        writer.append_data(image)
+    else:
+        with imageio.get_writer(path, fps=fps, mode='I') as writer:
+            for image in images:
+                writer.append_data(image)

src/utils/vis_proto.py

+# -*- coding: utf-8 -*-
+# Copyright 2025 BAAI. and/or its affiliates.
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+import sys
+import argparse
+import numpy as np
+from PIL import Image
+import io
+
+project_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+if project_root not in sys.path:
+    sys.path.insert(0, project_root)
+from src.proto import emu_pb as story_pb
+from src.utils.video_utils import plot_string, save_image_list_to_video
+
+def main():
+    parser = argparse.ArgumentParser(description='Visualize protobuf story files')
+    parser.add_argument('--input', '-i', required=True, help='Input protobuf file path')
+    parser.add_argument('--output', '-o', required=True, help='Output directory path')
+    parser.add_argument('--video', action='store_true', help='Generate video from protobuf content')
+    parser.add_argument('--fps', type=int, default=1, help='Frames per second for video (default: 1)')
+    args = parser.parse_args()
+    
+    input_path = args.input
+    output_path = args.output
+    
+    os.makedirs(output_path, exist_ok=True)
+    
+    with open(input_path, 'rb') as f:
+        story = story_pb.Story()
+        story.ParseFromString(f.read())
+
+    with open(f"{output_path}/000_question.txt", 'w') as f:
+        print(story.question, file=f)
+
+    if story.summary and story.summary.strip():
+        with open(f"{output_path}/000_global_cot.txt", 'w') as f:
+            print(story.summary, file=f)
+
+    idx = 1
+    
+    if len(story.reference_images) > 0:
+        for i in range(len(story.reference_images)):
+            with open(f"{output_path}/{i:03d}_reference_image.png", 'wb') as f:
+                f.write(story.reference_images[i].image.image_data)
+        idx = len(story.reference_images)
+        
+    for c in story.clips:
+        for s in c.segments:
+            with open(f"{output_path}/{idx:03d}_text.txt", 'w') as f:
+                print(s.asr, file=f)
+            for im_idx, im in enumerate(s.images):
+                with open(f"{output_path}/{idx:03d}_{im_idx:02d}_image.png", 'wb') as f:
+                    f.write(im.image.image_data)
+                if im.chain_of_thought and im.chain_of_thought.strip():
+                    with open(f"{output_path}/{idx:03d}_{im_idx:02d}_image_cot.txt", 'w') as f:
+                        print(im.chain_of_thought, file=f)
+            idx += 1
+    
+    if args.video:
+        video_images = []
+        target_size = None
+        
+        for ref_img_data in story.reference_images:
+            img = Image.open(io.BytesIO(ref_img_data.image.image_data))
+            img = img.convert('RGB')
+            if target_size is None:
+                target_size = img.size
+        
+        for c in story.clips:
+            for s in c.segments:
+                for im in s.images:
+                    img = Image.open(io.BytesIO(im.image.image_data))
+                    img = img.convert('RGB')
+                    if target_size is None:
+                        target_size = img.size
+        
+        if target_size is None:
+            target_size = (512, 512)
+        
+        if story.question and story.question.strip():
+            question_img = plot_string(story.question, image_size=(target_size[0], target_size[1]))
+            video_images.append(question_img)
+        
+        for img_array in story.reference_images:
+            img = Image.open(io.BytesIO(img_array.image.image_data))
+            img = img.convert('RGB')
+            if img.size != target_size:
+                img = img.resize(target_size, Image.Resampling.LANCZOS)
+            video_images.append(np.array(img))
+        
+        for c in story.clips:
+            for s in c.segments:
+                if s.asr and s.asr.strip():
+                    asr_img = plot_string(s.asr, image_size=(target_size[0], target_size[1]))
+                    video_images.append(asr_img)
+                
+                for im in s.images:
+                    img = Image.open(io.BytesIO(im.image.image_data))
+                    img = img.convert('RGB')
+                    if img.size != target_size:
+                        img = img.resize(target_size, Image.Resampling.LANCZOS)
+                    video_images.append(np.array(img))
+        
+        if video_images:
+            video_path = f"{output_path}/video.mp4"
+            save_image_list_to_video(video_images, video_path, fps=args.fps, quality='high')
+            print(f"Video saved to: {video_path}")
+
+if __name__ == "__main__":
+    main()

src/vision_tokenizer/__init__.py

+import os.path as osp
+
+from omegaconf import OmegaConf
+import torch
+
+from .ibq import IBQ
+
+
+def build_vision_tokenizer(type, model_path, device="cuda:0", **kwargs):
+    match type:
+        case "ibq":
+            cfg = OmegaConf.load(osp.join(model_path, "config.yaml"))
+            tokenizer = IBQ(**cfg)
+            ckpt = torch.load(osp.join(model_path, "model.ckpt"), map_location="cpu")
+            tokenizer.load_state_dict(ckpt)
+            tokenizer.eval().to(device)
+            return tokenizer
+        case _:
+            raise NotImplementedError(f"Unsupported vision tokenizer type: {type}")
+

src/vision_tokenizer/ibq.py

+# -*- coding:utf-8 -*-
+
+from torch import nn
+
+from .modules.diffusionmodules.model import Encoder, Decoder
+from .modules.vqvae.quantize import IndexPropagationQuantize
+
+class IBQ(nn.Module):
+
+    def __init__(
+        self,
+        ddconfig,
+        n_embed,
+        embed_dim,
+        beta=0.25,
+        use_entropy_loss=False,
+        cosine_similarity=False,
+        entropy_temperature=0.01,
+        sample_minimization_weight=1.0,
+        batch_maximization_weight=1.0,
+        **kwargs,
+    ):
+        super().__init__()
+
+        self.encoder = Encoder(**ddconfig)
+        self.decoder = Decoder(**ddconfig)
+
+        self.quantize = IndexPropagationQuantize(
+            n_embed,
+            embed_dim,
+            beta,
+            use_entropy_loss,
+            cosine_similarity=cosine_similarity,
+            entropy_temperature=entropy_temperature,
+            sample_minimization_weight=sample_minimization_weight,
+            batch_maximization_weight=batch_maximization_weight,
+        )
+
+        self.quant_conv = nn.Conv2d(ddconfig["z_channels"], embed_dim, 1)
+        self.post_quant_conv = nn.Conv2d(embed_dim, ddconfig["z_channels"], 1)
+
+    def encode(self, x):
+        h = self.encoder(x)
+        h = self.quant_conv(h)
+        quant, emb_loss, info = self.quantize(h)
+        return quant, emb_loss, info
+
+    def decode(self, quant, return_intermediate_feature=False):
+        quant = self.post_quant_conv(quant)
+        dec = self.decoder(quant, return_intermediate_feature=return_intermediate_feature)
+        return dec
+
+    def decode_code(self, code_b, shape=None):
+        # shape specifying (batch, height, width, channel)
+        quant_b = self.quantize.get_codebook_entry(code_b, shape=shape)
+        dec = self.decode(quant_b)
+        return dec
+
+    def forward(self, input, return_intermediate_feature=False):
+        quant, diff, _ = self.encode(input)
+        dec = self.decode(quant, return_intermediate_feature=return_intermediate_feature)
+        return dec, diff