diff --git a/extensions_built_in/diffusion_models/__init__.py b/extensions_built_in/diffusion_models/__init__.py index 30f4782..f30b4d2 100644 --- a/extensions_built_in/diffusion_models/__init__.py +++ b/extensions_built_in/diffusion_models/__init__.py @@ -11,6 +11,7 @@ from .ltx2 import LTX2Model, LTX23Model from .zeta_chroma import ZetaChromaModel from .ernie_image import ErnieImageModel from .nucleus_image import NucleusImageModel +from .hidream.hidream_o1_model import HidreamO1Model AI_TOOLKIT_MODELS = [ # put a list of models here @@ -36,4 +37,5 @@ AI_TOOLKIT_MODELS = [ ZetaChromaModel, ErnieImageModel, NucleusImageModel, + HidreamO1Model, ] diff --git a/extensions_built_in/diffusion_models/hidream/__init__.py b/extensions_built_in/diffusion_models/hidream/__init__.py index 32562a9..443d1b0 100644 --- a/extensions_built_in/diffusion_models/hidream/__init__.py +++ b/extensions_built_in/diffusion_models/hidream/__init__.py @@ -1,2 +1,3 @@ from .hidream_model import HidreamModel -from .hidream_e1_model import HidreamE1Model \ No newline at end of file +from .hidream_e1_model import HidreamE1Model +from .hidream_o1_model import HidreamO1Model diff --git a/extensions_built_in/diffusion_models/hidream/hidream_o1_model.py b/extensions_built_in/diffusion_models/hidream/hidream_o1_model.py new file mode 100644 index 0000000..28f60e3 --- /dev/null +++ b/extensions_built_in/diffusion_models/hidream/hidream_o1_model.py @@ -0,0 +1,544 @@ +import os +from typing import List, Optional + +import torch +import yaml +from toolkit.config_modules import GenerateImageConfig, ModelConfig +from toolkit.metadata import get_meta_for_safetensors +from toolkit.models.base_model import BaseModel +from toolkit.basic import flush +from toolkit.advanced_prompt_embeds import AdvancedPromptEmbeds +from toolkit.prompt_utils import PromptEmbeds +from toolkit.samplers.custom_flowmatch_sampler import ( + CustomFlowMatchEulerDiscreteScheduler, +) +from safetensors.torch import load_file, save_file +from toolkit.accelerator import unwrap_model +from optimum.quanto import freeze +from toolkit.util.quantize import quantize, get_qtype, quantize_model +from toolkit.memory_management import MemoryManager + +from transformers import AutoProcessor +from transformers.models.qwen3_vl.configuration_qwen3_vl import Qwen3VLConfig +from .src.hidream_o1.qwen3_vl_transformers import Qwen3VLForConditionalGeneration +from .src.hidream_o1.pipeline import HiDreamO1Pipeline, DEFAULT_NOISE_SCALE +from toolkit.models.FakeVAE import FakeVAE +from typing import TYPE_CHECKING +from .src.hidream_o1.model_config import model_config + +if TYPE_CHECKING: + from toolkit.data_transfer_object.data_loader import DataLoaderBatchDTO + +scheduler_config = { + "num_train_timesteps": 1000, + "shift": 3.0, + "use_dynamic_shifting": False, +} + +_GLOBAL_NOISE_SCALE = DEFAULT_NOISE_SCALE + + +class HidreamO1FlowmatchScheduler(CustomFlowMatchEulerDiscreteScheduler): + def __init__(self, *args, **kwargs): + self.noise_scale = kwargs.get("noise_scale", DEFAULT_NOISE_SCALE) + # remove noise_scale from kwargs so it doesn't get passed to the parent class + kwargs.pop("noise_scale", None) + super().__init__(*args, **kwargs) + + def add_noise( + self, + original_samples: torch.Tensor, + noise: torch.Tensor, + timesteps: torch.Tensor, + ) -> torch.Tensor: + t_01 = (timesteps / 1000).to(original_samples.device) + scaled_noise = noise * self.noise_scale + noisy_model_input = (1.0 - t_01) * original_samples + t_01 * scaled_noise + return noisy_model_input + + +def add_special_tokens(tokenizer): + """Attach the special-token shortcuts that the pipeline relies on.""" + tokenizer.boi_token = "<|boi_token|>" + tokenizer.bor_token = "<|bor_token|>" + tokenizer.eor_token = "<|eor_token|>" + tokenizer.bot_token = "<|bot_token|>" + tokenizer.tms_token = "<|tms_token|>" + + +def get_tokenizer(processor): + from transformers import PreTrainedTokenizerBase + + if isinstance(processor, PreTrainedTokenizerBase): + return processor + return processor.tokenizer + + +class FakeConfig: + pass + + +class FakeTextEncoder(torch.nn.Module): + def __init__(self, scaling_factor=1.0): + super().__init__() + self._dtype = torch.float32 + self._device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + self.config = FakeConfig() + self.config.scaling_factor = scaling_factor + + @property + def dtype(self): + return self._dtype + + @dtype.setter + def dtype(self, value): + self._dtype = value + + @property + def device(self): + return self._device + + @device.setter + def device(self, value): + self._device = value + + # mimic to from torch + def to(self, *args, **kwargs): + # pull out dtype and device if they exist + if "dtype" in kwargs: + self._dtype = kwargs["dtype"] + if "device" in kwargs: + self._device = kwargs["device"] + return super().to(*args, **kwargs) + + +class HidreamO1Model(BaseModel): + arch = "hidream_o1" + + def __init__( + self, + device, + model_config: ModelConfig, + dtype="bf16", + custom_pipeline=None, + noise_scheduler=None, + **kwargs, + ): + super().__init__( + device, model_config, dtype, custom_pipeline, noise_scheduler, **kwargs + ) + self.is_flow_matching = True + self.is_transformer = True + self.target_lora_modules = ["Qwen3VLForConditionalGeneration"] + self.noise_scale = self.model_config.model_kwargs.get( + "noise_scale", DEFAULT_NOISE_SCALE + ) + self.noise_scale_inference = self.model_config.model_kwargs.get( + "noise_scale_inference", self.noise_scale + ) + print(f"Using noise scale: {self.noise_scale}") + global _GLOBAL_NOISE_SCALE + _GLOBAL_NOISE_SCALE = self.noise_scale + self.is_comfy_weight = False # save as single file if true + + # static method to get the noise scheduler + @staticmethod + def get_train_scheduler(): + return HidreamO1FlowmatchScheduler( + **scheduler_config, noise_scale=_GLOBAL_NOISE_SCALE + ) + + def get_bucket_divisibility(self): + return 32 # patch size + + def load_model(self): + dtype = self.torch_dtype + self.print_and_status_update("Loading HidreamO1 model") + model_path = self.model_config.name_or_path + + self.print_and_status_update("Loading transformer") + + try: + processor = AutoProcessor.from_pretrained(model_path) + except Exception as e: + print( + f"Failed to load processor from model path {model_path}, trying original path. Error: {e}" + ) + processor_path = self.model_config.extras_name_or_path + if processor_path.endswith(".safetensors"): + processor_path = "HiDream-ai/HiDream-O1-Image" + processor = AutoProcessor.from_pretrained(processor_path) + + tokenizer = get_tokenizer(processor) + add_special_tokens(tokenizer) + + if model_path.endswith(".safetensors"): + self.is_comfy_weight = True + self.print_and_status_update( + "Model is in safetensors format, loading with safetensors" + ) + state_dict = load_file(model_path) + + for key, value in state_dict.items(): + state_dict[key] = value.to(dtype=dtype) + + # comfy ui is missing the lm head. It isnt used, but our model needs it for now + state_dict["lm_head.weight"] = torch.zeros( + 151936, 4096, dtype=torch.bfloat16, device="cpu" + ) + + # transformer.load_state_dict(state_dict, assign=True) + transformer = Qwen3VLForConditionalGeneration.from_pretrained( + None, + config=Qwen3VLConfig(**model_config), + state_dict=state_dict, + torch_dtype=self.torch_dtype, + ) + del state_dict # free memory + else: + transformer = Qwen3VLForConditionalGeneration.from_pretrained( + model_path, + torch_dtype=self.torch_dtype, + ) + flush() + if not self.model_config.low_vram: + transformer.to(self.device_torch) + + if self.model_config.quantize: + self.print_and_status_update("Quantizing Transformer") + quantize_model(self, transformer) + flush() + + if ( + self.model_config.layer_offloading + and self.model_config.layer_offloading_transformer_percent > 0 + ): + MemoryManager.attach( + transformer, + self.device_torch, + offload_percent=self.model_config.layer_offloading_transformer_percent, + ignore_modules=[], + ) + + flush() + + # move over to device now if low vram + if self.model_config.low_vram: + transformer.to(self.device_torch) + + # fake ones so the trainer doesnt break + vae = FakeVAE().to(self.device_torch, dtype=dtype) + text_encoder = FakeTextEncoder().to(self.device_torch, dtype=dtype) + + self.noise_scheduler = HidreamO1Model.get_train_scheduler() + + self.print_and_status_update("Making pipe") + + kwargs = {} + + pipe: HiDreamO1Pipeline = HiDreamO1Pipeline( + scheduler=self.noise_scheduler, + processor=processor, + model=None, + **kwargs, + ) + pipe.model = transformer + + self.print_and_status_update("Preparing Model") + + flush() + + # save it to the model class + self.vae = vae + self.text_encoder = text_encoder + self.tokenizer = processor + self.model = pipe.model + self.pipeline = pipe + self.print_and_status_update("Model Loaded") + + def get_generation_pipeline(self): + scheduler = HidreamO1Model.get_train_scheduler() + + pipe: HiDreamO1Pipeline = HiDreamO1Pipeline( + scheduler=scheduler, + processor=self.tokenizer, + model=None, + ) + pipe.model = self.transformer + + return pipe + + def encode_images(self, image_list: torch.Tensor, device=None, dtype=None): + if self.vae.device == torch.device("cpu"): + self.vae.to(self.device_torch) + if device is None: + device = self.vae_device_torch + if dtype is None: + dtype = self.vae_torch_dtype + + # not needed since there is not a latent space + return image_list.to(device, dtype=dtype) + + def decode_latents(self, latents: torch.Tensor, device=None, dtype=None): + if self.vae.device == torch.device("cpu"): + self.vae.to(self.device_torch) + if device is None: + device = self.vae_device_torch + if dtype is None: + dtype = self.vae_torch_dtype + + # not needed since there is not a latent space + return latents.to(device, dtype=dtype) + + def generate_single_image( + self, + pipeline: HiDreamO1Pipeline, + gen_config: GenerateImageConfig, + conditional_embeds: AdvancedPromptEmbeds, + unconditional_embeds: AdvancedPromptEmbeds, + generator: torch.Generator, + extra: dict, + ): + if self.model.device == torch.device("cpu"): + self.model.to(self.device_torch) + + sc = self.get_bucket_divisibility() + gen_config.width = int(gen_config.width // sc * sc) + gen_config.height = int(gen_config.height // sc * sc) + + img = pipeline( + # prompt=gen_config.prompt, + prompt_input_ids=conditional_embeds.text_embeds[0], + # negative_prompt=gen_config.negative_prompt, + negative_prompt_input_ids=unconditional_embeds.text_embeds[0], + height=gen_config.height, + width=gen_config.width, + num_inference_steps=gen_config.num_inference_steps, + guidance_scale=gen_config.guidance_scale, + generator=generator, + noise_scale=self.noise_scale_inference, + **extra, + ).images[0] + return img + + def get_noise_prediction( + self, + latent_model_input: torch.Tensor, + timestep: torch.Tensor, # 0 to 1000 scale + text_embeddings: AdvancedPromptEmbeds, + batch: "DataLoaderBatchDTO", + **kwargs, + ): + import einops + from .src.hidream_o1.pipeline import PATCH_SIZE, T_EPS + + if self.model.device == torch.device("cpu"): + self.model.to(self.device_torch) + + device = self.device_torch + in_dtype = latent_model_input.dtype + bs, _, h_pix, w_pix = latent_model_input.shape + h_patches = h_pix // PATCH_SIZE + w_patches = w_pix // PATCH_SIZE + + # (B, C, H, W) -> (B, H/p * W/p, C * p * p) + z = einops.rearrange( + latent_model_input, + "B C (H p1) (W p2) -> B (H W) (C p1 p2)", + p1=PATCH_SIZE, + p2=PATCH_SIZE, + ).to(device) + + model_config = self.model.config + pad_token_id = getattr(model_config, "pad_token_id", 0) or 0 + + with torch.no_grad(): + # Build per-sample conditioning, then left-pad the text portion so + # the boi/tms + vision-token suffix stays at the end of the + # sequence (the t2i layout assumes vision tokens are at the tail). + per_sample = [] + for b in range(bs): + tokens = text_embeddings.text_embeds[b] + if tokens.dim() == 1: + tokens = tokens.unsqueeze(0) + per_sample.append( + self.pipeline.build_conditioning_sample( + tokens.to(device), + h_pix, + w_pix, + ) + ) + + max_seq_len = max(s["input_ids"].shape[-1] for s in per_sample) + ids_l, pos_l, tt_l, vm_l, mask_l = [], [], [], [], [] + for s in per_sample: + ids = s["input_ids"].to(device) + pos = s["position_ids"].to(device) + tt = s["token_types"].to(device) + vm = s["vinput_mask"].to(device) + seq_len = ids.shape[-1] + pad_len = max_seq_len - seq_len + + if pad_len > 0: + ids = torch.cat( + [ + torch.full( + (1, pad_len), + pad_token_id, + dtype=ids.dtype, + device=device, + ), + ids, + ], + dim=-1, + ) + pos = torch.cat( + [ + torch.ones((3, 1, pad_len), dtype=pos.dtype, device=device), + pos, + ], + dim=-1, + ) + tt = torch.cat( + [ + torch.zeros((1, pad_len), dtype=tt.dtype, device=device), + tt, + ], + dim=-1, + ) + vm = torch.cat( + [ + torch.zeros((1, pad_len), dtype=vm.dtype, device=device), + vm, + ], + dim=-1, + ) + mask = torch.cat( + [ + torch.zeros((1, pad_len), dtype=torch.long, device=device), + torch.ones((1, seq_len), dtype=torch.long, device=device), + ], + dim=-1, + ) + else: + mask = torch.ones((1, seq_len), dtype=torch.long, device=device) + + ids_l.append(ids) + pos_l.append(pos) + tt_l.append(tt) + vm_l.append(vm) + mask_l.append(mask) + + input_ids = torch.cat(ids_l, dim=0) + position_ids = torch.cat(pos_l, dim=1) # (3, B, S) + token_types = torch.cat(tt_l, dim=0) + vinput_mask = torch.cat(vm_l, dim=0) + attention_mask = torch.cat(mask_l, dim=0) + + # Model wants timestep as denoising progress in (0, 1) where 1=clean. + t_pixeldit = (1.0 - timestep.float() / 1000.0).to(device) + + outputs = self.model( + input_ids=input_ids, + position_ids=position_ids, + attention_mask=attention_mask if bs > 1 else None, + vinputs=z, + timestep=t_pixeldit.reshape(-1), + token_types=token_types, + use_flash_attn=True, + ) + x_pred = outputs.x_pred # (B, S, C*p*p) over the full padded sequence + + # Pull the vision-token positions only. + vision_pred = torch.stack( + [x_pred[b][vinput_mask[b].bool()] for b in range(bs)], + dim=0, + ) # (B, image_len, C*p*p) + + x0_pred = einops.rearrange( + vision_pred, + "B (H W) (C p1 p2) -> B C (H p1) (W p2)", + H=h_patches, + W=w_patches, + p1=PATCH_SIZE, + p2=PATCH_SIZE, + ) + + # Model emits an x0-prediction; convert to flow-matching velocity + # (x_1 - x_0) so it matches the loss target from get_loss_target. + sigma = (timestep.float() / 1000.0).clamp_min(T_EPS).to(device) + while sigma.dim() < latent_model_input.dim(): + sigma = sigma.unsqueeze(-1) + pred = (latent_model_input.float().to(device) - x0_pred.float()) / sigma + return pred.to(in_dtype) + + def get_prompt_embeds(self, prompt: list) -> AdvancedPromptEmbeds: + if not isinstance(prompt, list): + prompt = [prompt] + # empty, we cannot use them with this omni model anyway, but will break trainer if they do not exist + token_list = [self.pipeline.encode_prompt(p) for p in prompt] + pe = AdvancedPromptEmbeds(text_embeds=token_list) + pe._frozen_dtype_keys = ["text_embeds"] + return pe + + def get_model_has_grad(self): + return False + + def get_te_has_grad(self): + return False + + def save_model(self, output_path, meta, save_dtype): + transformer: Qwen3VLForConditionalGeneration = unwrap_model(self.model) + if self.is_comfy_weight: + sd = transformer.state_dict() + save_dict = {} + for key, value in sd.items(): + if "lm_head.weight" in key: + continue # comfy checkpoint doesnt have the lm head, so skip it + save_dict[key] = value.clone().to("cpu", dtype=save_dtype) + + if not output_path.endswith(".safetensors"): + output_path += ".safetensors" + meta = get_meta_for_safetensors(meta, name=self.arch) + save_file(save_dict, output_path, metadata=meta) + else: + transformer.save_pretrained( + save_directory=output_path, + safe_serialization=True, + ) + + # save processor + self.tokenizer.save_pretrained(output_path) + + meta_path = os.path.join(output_path, "aitk_meta.yaml") + with open(meta_path, "w") as f: + yaml.dump(meta, f) + + def get_loss_target(self, *args, **kwargs): + noise = kwargs.get("noise") + batch = kwargs.get("batch") + noise_scale = self.noise_scale + return (noise * noise_scale - batch.latents).detach() + + def get_base_model_version(self): + return self.arch + + def get_transformer_block_names(self) -> Optional[List[str]]: + return ["layers"] + + def convert_lora_weights_before_save(self, state_dict): + new_sd = {} + for key, value in state_dict.items(): + new_key = key.replace("transformer.", "diffusion_model.") + new_key = new_key.replace(".model.", ".") + new_sd[new_key] = value + return new_sd + + def convert_lora_weights_before_load(self, state_dict): + new_sd = {} + for key, value in state_dict.items(): + new_key = key.replace("diffusion_model.", "transformer.model.") + # to load legacy keys + new_key = new_key.replace("transformer.model.model.", "transformer.model.") + new_sd[new_key] = value + return new_sd diff --git a/extensions_built_in/diffusion_models/hidream/src/hidream_o1/__init__.py b/extensions_built_in/diffusion_models/hidream/src/hidream_o1/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/extensions_built_in/diffusion_models/hidream/src/hidream_o1/model_config.py b/extensions_built_in/diffusion_models/hidream/src/hidream_o1/model_config.py new file mode 100644 index 0000000..ded2ff2 --- /dev/null +++ b/extensions_built_in/diffusion_models/hidream/src/hidream_o1/model_config.py @@ -0,0 +1,52 @@ +model_config = { + "architectures": ["Qwen3VLForConditionalGeneration"], + "image_token_id": 151655, + "model_type": "qwen3_vl", + "text_config": { + "attention_bias": False, + "attention_dropout": 0.0, + "bos_token_id": 151643, + "dtype": "bfloat16", + "eos_token_id": 151645, + "head_dim": 128, + "hidden_act": "silu", + "hidden_size": 4096, + "initializer_range": 0.02, + "intermediate_size": 12288, + "max_position_embeddings": 262144, + "model_type": "qwen3_vl_text", + "num_attention_heads": 32, + "num_hidden_layers": 36, + "num_key_value_heads": 8, + "rms_norm_eps": 1e-06, + "rope_scaling": { + "mrope_interleaved": True, + "mrope_section": [24, 20, 20], + "rope_type": "default", + }, + "rope_theta": 5000000, + "use_cache": True, + "vocab_size": 151936, + }, + "tie_word_embeddings": False, + "transformers_version": "4.57.0.dev0", + "video_token_id": 151656, + "vision_config": { + "deepstack_visual_indexes": [8, 16, 24], + "depth": 27, + "hidden_act": "gelu_pytorch_tanh", + "hidden_size": 1152, + "in_channels": 3, + "initializer_range": 0.02, + "intermediate_size": 4304, + "model_type": "qwen3_vl", + "num_heads": 16, + "num_position_embeddings": 2304, + "out_hidden_size": 4096, + "patch_size": 16, + "spatial_merge_size": 2, + "temporal_patch_size": 2, + }, + "vision_end_token_id": 151653, + "vision_start_token_id": 151652, +} diff --git a/extensions_built_in/diffusion_models/hidream/src/hidream_o1/pipeline.py b/extensions_built_in/diffusion_models/hidream/src/hidream_o1/pipeline.py new file mode 100644 index 0000000..8e24da3 --- /dev/null +++ b/extensions_built_in/diffusion_models/hidream/src/hidream_o1/pipeline.py @@ -0,0 +1,455 @@ +from typing import List, Optional, Union + +import einops +import numpy as np +import torch +from PIL import Image +import torchvision.transforms.v2 as transforms + +from diffusers import DiffusionPipeline, FlowMatchEulerDiscreteScheduler +from diffusers.utils import BaseOutput +from dataclasses import dataclass + + +TIMESTEP_TOKEN_NUM = 1 +DEFAULT_NOISE_SCALE = 8.0 +T_EPS = 0.001 +PATCH_SIZE = 32 + +TENSOR_TRANSFORM = transforms.Compose( + [ + transforms.ToImage(), + transforms.ToDtype(torch.float32, scale=True), + transforms.Normalize([0.5], [0.5]), + ] +) + + +def round_to_patch(dim: int, patch: int = PATCH_SIZE) -> int: + return max(patch, int(dim // patch * patch)) + + +def _get_rope_index_t2i( + spatial_merge_size: int, + image_token_id: int, + video_token_id: int, + vision_start_token_id: int, + input_ids: torch.LongTensor, + image_grid_thw: torch.LongTensor, + skip_vision_start_token: List[int], + fix_point: int = 4096, +): + """Compute mrope position ids for the t2i case used by HiDream-O1.""" + attention_mask = torch.ones_like(input_ids) + position_ids = torch.ones( + 3, + input_ids.shape[0], + input_ids.shape[1], + dtype=input_ids.dtype, + device=input_ids.device, + ) + + for i, ids_row in enumerate(input_ids): + ids_row = ids_row[attention_mask[i] == 1] + vision_start_indices = torch.argwhere(ids_row == vision_start_token_id).squeeze( + 1 + ) + vision_tokens = ids_row[vision_start_indices + 1] + image_nums = (vision_tokens == image_token_id).sum().item() + video_nums = (vision_tokens == video_token_id).sum().item() + input_tokens = ids_row.tolist() + + llm_pos_ids_list = [] + st = 0 + image_index = 0 + video_index = 0 + remain_images, remain_videos = image_nums, video_nums + local_fix_point = fix_point + + for _ in range(image_nums + video_nums): + ed_image = ( + input_tokens.index(image_token_id, st) + if (image_token_id in input_tokens and remain_images > 0) + else len(input_tokens) + 1 + ) + ed_video = ( + input_tokens.index(video_token_id, st) + if (video_token_id in input_tokens and remain_videos > 0) + else len(input_tokens) + 1 + ) + if ed_image < ed_video: + t, h, w = image_grid_thw[image_index].tolist() + image_index += 1 + remain_images -= 1 + ed = ed_image + else: + t, h, w = image_grid_thw[video_index].tolist() + video_index += 1 + remain_videos -= 1 + ed = ed_video + + llm_grid_t = t + llm_grid_h = h // spatial_merge_size + llm_grid_w = w // spatial_merge_size + + text_len = ed - st - skip_vision_start_token[image_index - 1] + text_len = max(0, text_len) + + st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 + llm_pos_ids_list.append( + torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx + ) + + t_index = ( + torch.arange(llm_grid_t) + .view(-1, 1) + .expand(-1, llm_grid_h * llm_grid_w) + .flatten() + ) + h_index = ( + torch.arange(llm_grid_h) + .view(1, -1, 1) + .expand(llm_grid_t, -1, llm_grid_w) + .flatten() + ) + w_index = ( + torch.arange(llm_grid_w) + .view(1, 1, -1) + .expand(llm_grid_t, llm_grid_h, -1) + .flatten() + ) + + if skip_vision_start_token[image_index - 1]: + if local_fix_point > 0: + local_fix_point = local_fix_point - st_idx + llm_pos_ids_list.append( + torch.stack([t_index, h_index, w_index]) + local_fix_point + st_idx + ) + local_fix_point = 0 + else: + llm_pos_ids_list.append( + torch.stack([t_index, h_index, w_index]) + text_len + st_idx + ) + st = ed + llm_grid_t * llm_grid_h * llm_grid_w + + if st < len(input_tokens): + st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 + text_len = len(input_tokens) - st + llm_pos_ids_list.append( + torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx + ) + + llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) + position_ids[..., i, attention_mask[i] == 1] = llm_positions.to( + position_ids.device + ) + + return position_ids + + +def _build_t2i_sample_from_input_ids( + input_ids: torch.Tensor, + height: int, + width: int, + model_config, + attention_mask: Optional[torch.Tensor] = None, +): + """Build the full conditioning sample (position_ids/token_types/vinput_mask) + around an already-tokenized prompt.""" + image_token_id = model_config.image_token_id + video_token_id = model_config.video_token_id + vision_start_token_id = model_config.vision_start_token_id + image_len = (height // PATCH_SIZE) * (width // PATCH_SIZE) + + if input_ids.dim() == 1: + input_ids = input_ids.unsqueeze(0) + + image_grid_thw = torch.tensor( + [1, height // PATCH_SIZE, width // PATCH_SIZE], dtype=torch.int64 + ).unsqueeze(0) + + vision_tokens = ( + torch.zeros((1, image_len), dtype=input_ids.dtype, device=input_ids.device) + + image_token_id + ) + vision_tokens[0, 0] = vision_start_token_id + input_ids_pad = torch.cat([input_ids, vision_tokens], dim=-1) + + position_ids = _get_rope_index_t2i( + spatial_merge_size=1, + image_token_id=image_token_id, + video_token_id=video_token_id, + vision_start_token_id=vision_start_token_id, + input_ids=input_ids_pad, + image_grid_thw=image_grid_thw, + skip_vision_start_token=[1], + ) + + txt_seq_len = input_ids.shape[-1] + all_seq_len = position_ids.shape[-1] + + token_types = torch.zeros((1, all_seq_len), dtype=input_ids.dtype) + bgn = txt_seq_len - TIMESTEP_TOKEN_NUM + token_types[0, bgn : bgn + image_len + TIMESTEP_TOKEN_NUM] = 1 + token_types[0, txt_seq_len - TIMESTEP_TOKEN_NUM : txt_seq_len] = 3 + + vinput_mask = token_types == 1 + token_types_bin = (token_types > 0).to(token_types.dtype) + + sample = { + "input_ids": input_ids, + "position_ids": position_ids, + "token_types": token_types_bin, + "vinput_mask": vinput_mask, + } + if attention_mask is not None: + if attention_mask.dim() == 1: + attention_mask = attention_mask.unsqueeze(0) + sample["attention_mask"] = attention_mask + return sample + + +@dataclass +class HiDreamO1PipelineOutput(BaseOutput): + images: List[Image.Image] + + +class HiDreamO1Pipeline(DiffusionPipeline): + """ + Diffusers-style inference pipeline for HiDream-O1 (base model). + + HiDream-O1 is a unified text/vision/diffusion model with no VAE — the + transformer directly predicts image patches in pixel space. This pipeline + keeps only the components needed for text-to-image inference. + """ + + model_cpu_offload_seq = "model" + + def __init__( + self, + model, + processor, + scheduler: FlowMatchEulerDiscreteScheduler, + ): + super().__init__() + self.register_modules(model=model, processor=processor, scheduler=scheduler) + + @property + def tokenizer(self): + return ( + self.processor.tokenizer + if hasattr(self.processor, "tokenizer") + else self.processor + ) + + def _snap_resolution(self, width: int, height: int): + w, h = round_to_patch(width), round_to_patch(height) + if (w, h) != (width, height): + print(f"[hidream-o1] Resolution rounded from {width}x{height} to {w}x{h}") + return w, h + + def build_conditioning_sample( + self, + input_ids: torch.Tensor, + height: int, + width: int, + attention_mask: Optional[torch.Tensor] = None, + ): + """Build the per-sample conditioning dict (input_ids, position_ids, + token_types, vinput_mask) around already-tokenized text. Useful when + a training loop needs to batch samples manually.""" + return _build_t2i_sample_from_input_ids( + input_ids, + height, + width, + self.model.config, + attention_mask=attention_mask, + ) + + def encode_prompt(self, prompt: str) -> torch.Tensor: + """Apply the chat template + boi/tms suffix and tokenize. + Returns input_ids of shape (1, seq_len). Use these to precompute and + pass back into __call__ via `prompt_input_ids` / `negative_prompt_input_ids`.""" + tokenizer = self.tokenizer + boi_token = getattr(tokenizer, "boi_token", "<|boi_token|>") + tms_token = getattr(tokenizer, "tms_token", "<|tms_token|>") + + messages = [{"role": "user", "content": prompt}] + template_caption = ( + self.processor.apply_chat_template( + messages, tokenize=False, add_generation_prompt=True + ) + + boi_token + + tms_token * TIMESTEP_TOKEN_NUM + ) + return tokenizer.encode( + template_caption, return_tensors="pt", add_special_tokens=False + ) + + @torch.no_grad() + def __call__( + self, + prompt: Optional[Union[str, List[str]]] = None, + negative_prompt: Optional[Union[str, List[str]]] = " ", + prompt_input_ids: Optional[torch.Tensor] = None, + negative_prompt_input_ids: Optional[torch.Tensor] = None, + prompt_attention_mask: Optional[torch.Tensor] = None, + negative_prompt_attention_mask: Optional[torch.Tensor] = None, + height: int = 1440, + width: int = 2560, + num_inference_steps: int = 50, + guidance_scale: float = 5.0, + shift: float = 3.0, + generator: Optional[torch.Generator] = None, + seed: Optional[int] = None, + noise_scale: float = None, + output_type: str = "pil", + return_dict: bool = True, + ): + if noise_scale is None: + noise_scale = DEFAULT_NOISE_SCALE + if prompt is None and prompt_input_ids is None: + raise ValueError("Provide either `prompt` or `prompt_input_ids`.") + + def _unwrap_str(x): + if isinstance(x, list): + if len(x) != 1: + raise ValueError( + "HiDreamO1Pipeline currently supports batch size 1." + ) + return x[0] + return x + + prompt = _unwrap_str(prompt) + negative_prompt = _unwrap_str(negative_prompt) + + device = self._execution_device + dtype = torch.bfloat16 + model_config = self.model.config + + width, height = self._snap_resolution(width, height) + h_patches = height // PATCH_SIZE + w_patches = width // PATCH_SIZE + + do_cfg = guidance_scale > 1.0 + + if prompt_input_ids is None: + prompt_input_ids = self.encode_prompt(prompt) + if do_cfg and negative_prompt_input_ids is None: + if negative_prompt is None: + negative_prompt = " " + negative_prompt_input_ids = self.encode_prompt(negative_prompt) + + cond_sample = _build_t2i_sample_from_input_ids( + prompt_input_ids, + height, + width, + model_config, + attention_mask=prompt_attention_mask, + ) + uncond_sample = ( + _build_t2i_sample_from_input_ids( + negative_prompt_input_ids, + height, + width, + model_config, + attention_mask=negative_prompt_attention_mask, + ) + if do_cfg + else None + ) + + def _to_device(s): + return { + k: (v.to(device) if torch.is_tensor(v) else v) for k, v in s.items() + } + + cond_sample = _to_device(cond_sample) + if uncond_sample is not None: + uncond_sample = _to_device(uncond_sample) + + if generator is None: + if seed is None: + seed = 0 + generator = torch.Generator(device="cpu").manual_seed(seed + 1) + torch.manual_seed(seed + 1) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed + 1) + + noise = noise_scale * torch.randn( + (1, 3, height, width), generator=generator + ).to(device, dtype) + z = einops.rearrange( + noise, + "B C (H p1) (W p2) -> B (H W) (C p1 p2)", + p1=PATCH_SIZE, + p2=PATCH_SIZE, + ) + + if shift is not None and hasattr(self.scheduler, "set_shift"): + self.scheduler.set_shift(shift) + self.scheduler.set_timesteps(num_inference_steps, device=device) + + timesteps = self.scheduler.timesteps + + def _forward_once(sample, z_in, t_pixeldit): + with torch.autocast(device.type, dtype=dtype): + kwargs = { + "input_ids": sample["input_ids"], + "position_ids": sample["position_ids"], + "vinputs": z_in, + "timestep": t_pixeldit.reshape(-1).to(device), + "token_types": sample["token_types"], + "use_flash_attn": True, + } + if "attention_mask" in sample: + kwargs["attention_mask"] = sample["attention_mask"] + outputs = self.model(**kwargs) + x_pred = outputs.x_pred + return x_pred[0, sample["vinput_mask"][0]].unsqueeze(0) + + for step_t in self.progress_bar(timesteps): + t_pixeldit = 1.0 - step_t.float() / 1000.0 + sigma = (step_t.float() / 1000.0).to(dtype=torch.float32).clamp_min(T_EPS) + + x_pred_cond = _forward_once(cond_sample, z.clone(), t_pixeldit) + v_cond = (x_pred_cond.float() - z.float()) / sigma + + if do_cfg: + x_pred_uncond = _forward_once(uncond_sample, z.clone(), t_pixeldit) + v_uncond = (x_pred_uncond.float() - z.float()) / sigma + v_guided = v_uncond + guidance_scale * (v_cond - v_uncond) + else: + v_guided = v_cond + + model_output = -v_guided + z = self.scheduler.step( + model_output.float(), + step_t.to(dtype=torch.float32), + z.float(), + return_dict=False, + )[0].to(dtype) + + img = (z + 1) / 2 + img = einops.rearrange( + img.cpu().float(), + "B (H W) (C p1 p2) -> B C (H p1) (W p2)", + H=h_patches, + W=w_patches, + p1=PATCH_SIZE, + p2=PATCH_SIZE, + ) + + if output_type == "pt": + images = img.clamp(0, 1) + elif output_type == "np": + images = np.clip(img.numpy().transpose(0, 2, 3, 1), 0, 1) + else: + arr = np.round( + np.clip(img[0].numpy().transpose(1, 2, 0) * 255, 0, 255) + ).astype(np.uint8) + images = [Image.fromarray(arr).convert("RGB")] + + if not return_dict: + return (images,) + return HiDreamO1PipelineOutput(images=images) diff --git a/extensions_built_in/diffusion_models/hidream/src/hidream_o1/qwen3_vl_transformers.py b/extensions_built_in/diffusion_models/hidream/src/hidream_o1/qwen3_vl_transformers.py new file mode 100644 index 0000000..4172d21 --- /dev/null +++ b/extensions_built_in/diffusion_models/hidream/src/hidream_o1/qwen3_vl_transformers.py @@ -0,0 +1,2465 @@ +# Copyright 2025 The Qwen Team and The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import os +from dataclasses import dataclass +from typing import Any, Callable, Optional, Union +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F + +USE_BF16_ROPE = os.environ.get("USE_BF16_ROPE", "0") +from transformers.activations import ACT2FN +from transformers.cache_utils import Cache, DynamicCache +from transformers.generation import GenerationMixin +from transformers.integrations import use_kernel_forward_from_hub +from transformers.masking_utils import create_causal_mask +from transformers.modeling_flash_attention_utils import FlashAttentionKwargs +from transformers.modeling_layers import GradientCheckpointingLayer +from transformers.modeling_outputs import BaseModelOutputWithPast, ModelOutput +from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update +from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from transformers.processing_utils import Unpack +from transformers.utils import ( + TransformersKwargs, + auto_docstring, + is_torchdynamo_compiling, +) +from transformers.utils.deprecation import deprecate_kwarg +from transformers.utils.generic import merge_with_config_defaults +from transformers.models.qwen3_vl.configuration_qwen3_vl import ( + Qwen3VLConfig, + Qwen3VLTextConfig, + Qwen3VLVisionConfig, +) + + +class Qwen3VLVisionMLP(nn.Module): + def __init__(self, config): + super().__init__() + self.hidden_size = config.hidden_size + self.intermediate_size = config.intermediate_size + self.linear_fc1 = nn.Linear(self.hidden_size, self.intermediate_size, bias=True) + self.linear_fc2 = nn.Linear(self.intermediate_size, self.hidden_size, bias=True) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, hidden_state): + return self.linear_fc2(self.act_fn(self.linear_fc1(hidden_state))) + + +class Qwen3VLVisionPatchEmbed(nn.Module): + def __init__(self, config) -> None: + super().__init__() + self.patch_size = config.patch_size + self.temporal_patch_size = config.temporal_patch_size + self.in_channels = config.in_channels + self.embed_dim = config.hidden_size + + kernel_size = [self.temporal_patch_size, self.patch_size, self.patch_size] + self.proj = nn.Conv3d( + self.in_channels, + self.embed_dim, + kernel_size=kernel_size, + stride=kernel_size, + bias=True, + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + target_dtype = self.proj.weight.dtype + hidden_states = hidden_states.view( + -1, + self.in_channels, + self.temporal_patch_size, + self.patch_size, + self.patch_size, + ) + hidden_states = self.proj(hidden_states.to(dtype=target_dtype)).view( + -1, self.embed_dim + ) + return hidden_states + + +class Qwen3VLVisionRotaryEmbedding(nn.Module): + inv_freq: torch.Tensor # fix linting for `register_buffer` + + def __init__(self, dim: int, theta: float = 10000.0) -> None: + super().__init__() + inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float) / dim)) + self.register_buffer("inv_freq", inv_freq, persistent=False) + + def forward(self, seqlen: int) -> torch.Tensor: + seq = torch.arange( + seqlen, device=self.inv_freq.device, dtype=self.inv_freq.dtype + ) + freqs = torch.outer(seq, self.inv_freq) + return freqs + + +class Qwen3VLVisionPatchMerger(nn.Module): + def __init__(self, config: Qwen3VLVisionConfig, use_postshuffle_norm=False) -> None: + super().__init__() + self.hidden_size = config.hidden_size * (config.spatial_merge_size**2) + self.use_postshuffle_norm = use_postshuffle_norm + self.norm = nn.LayerNorm( + self.hidden_size if use_postshuffle_norm else config.hidden_size, eps=1e-6 + ) + self.linear_fc1 = nn.Linear(self.hidden_size, self.hidden_size) + self.act_fn = nn.GELU() + self.linear_fc2 = nn.Linear(self.hidden_size, config.out_hidden_size) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.norm( + x.view(-1, self.hidden_size) if self.use_postshuffle_norm else x + ).view(-1, self.hidden_size) + x = self.linear_fc2(self.act_fn(self.linear_fc1(x))) + return x + + +def rotate_half(x): + """Rotates half the hidden dims of the input.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def apply_rotary_pos_emb_vision( + q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + orig_q_dtype = q.dtype + orig_k_dtype = k.dtype + q, k = q.float(), k.float() + cos, sin = cos.unsqueeze(-2).float(), sin.unsqueeze(-2).float() + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + q_embed = q_embed.to(orig_q_dtype) + k_embed = k_embed.to(orig_k_dtype) + return q_embed, k_embed + + +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand( + batch, num_key_value_heads, n_rep, slen, head_dim + ) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +def eager_attention_forward( + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: Optional[torch.Tensor], + scaling: float, + dropout: float = 0.0, + **kwargs: Unpack[TransformersKwargs], +): + key_states = repeat_kv(key, module.num_key_value_groups) + value_states = repeat_kv(value, module.num_key_value_groups) + + attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling + if attention_mask is not None: + causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] + attn_weights = attn_weights + causal_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to( + query.dtype + ) + attn_weights = nn.functional.dropout( + attn_weights, p=dropout, training=module.training + ) + attn_output = torch.matmul(attn_weights, value_states) + attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, attn_weights + + +class Qwen3VLVisionAttention(nn.Module): + def __init__(self, config: Qwen3VLVisionConfig) -> None: + super().__init__() + self.dim = config.hidden_size + self.num_heads = config.num_heads + self.head_dim = self.dim // self.num_heads + self.num_key_value_groups = 1 # needed for eager attention + self.qkv = nn.Linear(self.dim, self.dim * 3, bias=True) + self.proj = nn.Linear(self.dim, self.dim) + self.scaling = self.head_dim**-0.5 + self.config = config + self.attention_dropout = 0.0 + self.is_causal = False + + def forward( + self, + hidden_states: torch.Tensor, + cu_seqlens: torch.Tensor, + rotary_pos_emb: Optional[torch.Tensor] = None, + position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, + **kwargs, + ) -> torch.Tensor: + seq_length = hidden_states.shape[0] + query_states, key_states, value_states = ( + self.qkv(hidden_states) + .reshape(seq_length, 3, self.num_heads, -1) + .permute(1, 0, 2, 3) + .unbind(0) + ) + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb_vision( + query_states, key_states, cos, sin + ) + + query_states = query_states.transpose(0, 1).unsqueeze(0) + key_states = key_states.transpose(0, 1).unsqueeze(0) + value_states = value_states.transpose(0, 1).unsqueeze(0) + + attention_interface: Callable = eager_attention_forward + if self.config._attn_implementation != "eager": + attention_interface = ALL_ATTENTION_FUNCTIONS[ + self.config._attn_implementation + ] + + if self.config._attn_implementation == "flash_attention_2": + # Flash Attention 2: Use cu_seqlens for variable length attention + max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max() + attn_output, _ = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask=None, + scaling=self.scaling, + dropout=0.0 if not self.training else self.attention_dropout, + cu_seq_lens_q=cu_seqlens, + cu_seq_lens_k=cu_seqlens, + max_length_q=max_seqlen, + max_length_k=max_seqlen, + is_causal=False, + **kwargs, + ) + else: + # Other implementations: Process each chunk separately + lengths = cu_seqlens[1:] - cu_seqlens[:-1] + splits = [ + torch.split(tensor, lengths.tolist(), dim=2) + for tensor in (query_states, key_states, value_states) + ] + + attn_outputs = [ + attention_interface( + self, + q, + k, + v, + attention_mask=None, + scaling=self.scaling, + dropout=0.0 if not self.training else self.attention_dropout, + is_causal=False, + **kwargs, + )[0] + for q, k, v in zip(*splits) + ] + attn_output = torch.cat(attn_outputs, dim=1) + + attn_output = attn_output.reshape(seq_length, -1).contiguous() + attn_output = self.proj(attn_output) + return attn_output + + +class Qwen3VLVisionBlock(GradientCheckpointingLayer): + def __init__(self, config, attn_implementation: str = "sdpa") -> None: + super().__init__() + self.norm1 = nn.LayerNorm(config.hidden_size, eps=1e-6) + self.norm2 = nn.LayerNorm(config.hidden_size, eps=1e-6) + self.attn = Qwen3VLVisionAttention(config=config) + self.mlp = Qwen3VLVisionMLP(config=config) + + def forward( + self, + hidden_states: torch.Tensor, + cu_seqlens: torch.Tensor, + rotary_pos_emb: Optional[torch.Tensor] = None, + position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, + **kwargs, + ) -> torch.Tensor: + hidden_states = hidden_states + self.attn( + self.norm1(hidden_states), + cu_seqlens=cu_seqlens, + rotary_pos_emb=rotary_pos_emb, + position_embeddings=position_embeddings, + **kwargs, + ) + hidden_states = hidden_states + self.mlp(self.norm2(hidden_states)) + return hidden_states + + +class Qwen3VLTextRotaryEmbedding(nn.Module): + inv_freq: torch.Tensor # fix linting for `register_buffer` + + def __init__(self, config: Qwen3VLTextConfig, device=None): + super().__init__() + if hasattr(config, "rope_scaling") and config.rope_scaling is not None: + self.rope_type = config.rope_scaling.get("rope_type", "default") + else: + self.rope_type = "default" + self.max_seq_len_cached = config.max_position_embeddings + self.original_max_seq_len = config.max_position_embeddings + + self.config = config + # Newer transformers removed "default" from ROPE_INIT_FUNCTIONS in + # favor of a static method on the rotary module. Match that. + if self.rope_type == "default": + self.rope_init_fn = self.compute_default_rope_parameters + else: + self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] + + inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device) + self.register_buffer("inv_freq", inv_freq, persistent=False) + self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False) + + self.mrope_section = config.rope_scaling.get("mrope_section", [24, 20, 20]) + + @staticmethod + def compute_default_rope_parameters( + config=None, + device=None, + seq_len: Optional[int] = None, + ) -> tuple[torch.Tensor, float]: + """Standard RoPE inv_freq. Mirrors upstream Qwen3VL since newer + transformers removed "default" from ROPE_INIT_FUNCTIONS.""" + rope_params = getattr(config, "rope_parameters", None) + base = ( + rope_params["rope_theta"] + if rope_params and "rope_theta" in rope_params + else getattr(config, "rope_theta", 10000.0) + ) + dim = ( + getattr(config, "head_dim", None) + or config.hidden_size // config.num_attention_heads + ) + inv_freq = 1.0 / ( + base + ** ( + torch.arange(0, dim, 2, dtype=torch.int64).to( + device=device, dtype=torch.float + ) + / dim + ) + ) + return inv_freq, 1.0 + + def apply_interleaved_mrope(self, freqs, mrope_section): + """Apply interleaved MRoPE to 3D rotary embeddings. + Reorganizes frequency layout from chunked [TTT...HHH...WWW] to + interleaved [THTHWHTHW...TT], preserving frequency continuity. + args: + x: (3, bs, seq_len, head_dim // 2) + mrope_section: (3,) + returns: + x_t: (bs, seq_len, head_dim // 2) + """ + freqs_t = freqs[0] # just overwrite the first dimension T + for dim, offset in enumerate((1, 2), start=1): # H, W + length = mrope_section[dim] * 3 + idx = slice(offset, length, 3) + freqs_t[..., idx] = freqs[dim, ..., idx] + return freqs_t + + @torch.no_grad() + @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope) + def forward(self, x, position_ids): + # In contrast to other models, Qwen3VL has different position ids for the grids + # So we expand the inv_freq to shape (3, ...) + if position_ids.ndim == 2: + position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1) + if USE_BF16_ROPE == "1": + inv_freq_expanded = ( + self.inv_freq[None, None, :, None] + .float() + .to(device=x.device) + .expand(3, position_ids.shape[1], -1, 1) + ) + else: + inv_freq_expanded = ( + self.original_inv_freq[None, None, :, None] + .float() + .to(device=x.device) + .expand(3, position_ids.shape[1], -1, 1) + ) + # inv_freq_expanded = self.inv_freq[None, None, :, None].float().to(device=x.device).expand(3, position_ids.shape[1], -1, 1) + position_ids_expanded = position_ids[ + :, :, None, : + ].float() # shape (3, bs, 1, positions) + + device_type = ( + x.device.type + if isinstance(x.device.type, str) and x.device.type != "mps" + else "cpu" + ) + with torch.autocast(device_type=device_type, enabled=False): # Force float32 + freqs = ( + inv_freq_expanded.float() @ position_ids_expanded.float() + ).transpose(2, 3) + freqs = self.apply_interleaved_mrope(freqs, self.mrope_section) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() * self.attention_scaling + sin = emb.sin() * self.attention_scaling + + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + +@use_kernel_forward_from_hub("RMSNorm") +class Qwen3VLTextRMSNorm(nn.Module): + def __init__(self, hidden_size, eps: float = 1e-6) -> None: + """ + Qwen3VLTextRMSNorm is equivalent to T5LayerNorm + """ + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + return self.weight * hidden_states.to(input_dtype) + + def extra_repr(self): + return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" + + +def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): + """Applies Rotary Position Embedding to the query and key tensors. + + Args: + q (`torch.Tensor`): The query tensor. + k (`torch.Tensor`): The key tensor. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + position_ids (`torch.Tensor`, *optional*): + Deprecated and unused. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + +class Qwen3VLTextAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: Qwen3VLTextConfig, layer_idx: int): + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.head_dim = getattr( + config, "head_dim", config.hidden_size // config.num_attention_heads + ) + self.num_key_value_groups = ( + config.num_attention_heads // config.num_key_value_heads + ) + self.scaling = self.head_dim**-0.5 + self.attention_dropout = config.attention_dropout + self.is_causal = True + + self.q_proj = nn.Linear( + config.hidden_size, + config.num_attention_heads * self.head_dim, + bias=config.attention_bias, + ) + self.k_proj = nn.Linear( + config.hidden_size, + config.num_key_value_heads * self.head_dim, + bias=config.attention_bias, + ) + self.v_proj = nn.Linear( + config.hidden_size, + config.num_key_value_heads * self.head_dim, + bias=config.attention_bias, + ) + self.o_proj = nn.Linear( + config.num_attention_heads * self.head_dim, + config.hidden_size, + bias=config.attention_bias, + ) + self.q_norm = Qwen3VLTextRMSNorm( + self.head_dim, eps=config.rms_norm_eps + ) # unlike olmo, only on the head dim! + self.k_norm = Qwen3VLTextRMSNorm( + self.head_dim, eps=config.rms_norm_eps + ) # thus post q_norm does not need reshape + + @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58") + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + attention_mask: Optional[torch.Tensor], + past_key_values: Optional[Cache] = None, + cache_position: Optional[torch.LongTensor] = None, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + + query_states = self.q_norm( + self.q_proj(hidden_states).view(hidden_shape) + ).transpose(1, 2) + key_states = self.k_norm( + self.k_proj(hidden_states).view(hidden_shape) + ).transpose(1, 2) + value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) + + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb( + query_states, key_states, cos, sin + ) + + if past_key_values is not None: + # sin and cos are specific to RoPE models; cache_position needed for the static cache + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} + key_states, value_states = past_key_values.update( + key_states, value_states, self.layer_idx, cache_kwargs + ) + + attention_interface: Callable = eager_attention_forward + if self.config._attn_implementation != "eager": + attention_interface = ALL_ATTENTION_FUNCTIONS[ + self.config._attn_implementation + ] + + attn_output, attn_weights = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + **kwargs, + ) + + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights + + +class Qwen3VLTextMLP(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.intermediate_size = config.intermediate_size + self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, x): + down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + return down_proj + + +class Qwen3VLTextDecoderLayer(GradientCheckpointingLayer): + def __init__(self, config: Qwen3VLTextConfig, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + + self.self_attn = Qwen3VLTextAttention(config=config, layer_idx=layer_idx) + + self.mlp = Qwen3VLTextMLP(config) + self.input_layernorm = Qwen3VLTextRMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.post_attention_layernorm = Qwen3VLTextRMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58") + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Cache] = None, + use_cache: Optional[bool] = False, + cache_position: Optional[torch.LongTensor] = None, + **kwargs: Unpack[TransformersKwargs], + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + # Self Attention + hidden_states, _ = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + cache_position=cache_position, + position_embeddings=position_embeddings, + **kwargs, + ) + hidden_states = residual + hidden_states + + # Fully Connected + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + return hidden_states + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for Llava outputs, with hidden states and attentions. + """ +) +class Qwen3VLModelOutputWithPast(ModelOutput): + r""" + past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache). + + Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see + `past_key_values` input) to speed up sequential decoding. + rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*): + The rope index difference between sequence length and multimodal rope. + """ + + last_hidden_state: Optional[torch.FloatTensor] = None + past_key_values: Optional[Cache] = None + hidden_states: Optional[tuple[torch.FloatTensor]] = None + attentions: Optional[tuple[torch.FloatTensor]] = None + rope_deltas: Optional[torch.LongTensor] = None + x_pred: Optional[torch.FloatTensor] = None + mid_results: Optional[list] = None + + +@auto_docstring +class Qwen3VLPreTrainedModel(PreTrainedModel): + config: Qwen3VLConfig + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["Qwen3VLTextDecoderLayer", "Qwen3VLVisionBlock"] + _skip_keys_device_placement = "past_key_values" + _supports_flash_attn = True + _supports_sdpa = True + + _can_compile_fullgraph = True + _supports_attention_backend = True + _can_record_outputs = { + "hidden_states": Qwen3VLTextDecoderLayer, + "attentions": Qwen3VLTextAttention, + } + + +class Qwen3VLVisionModel(Qwen3VLPreTrainedModel): + config: Qwen3VLVisionConfig + _no_split_modules = ["Qwen3VLVisionBlock"] + + def __init__(self, config, *inputs, **kwargs) -> None: + super().__init__(config, *inputs, **kwargs) + self.spatial_merge_size = config.spatial_merge_size + self.patch_size = config.patch_size + self.spatial_merge_unit = self.spatial_merge_size * self.spatial_merge_size + + self.patch_embed = Qwen3VLVisionPatchEmbed( + config=config, + ) + + self.pos_embed = nn.Embedding( + config.num_position_embeddings, config.hidden_size + ) + self.num_grid_per_side = int(config.num_position_embeddings**0.5) + + head_dim = config.hidden_size // config.num_heads + self.rotary_pos_emb = Qwen3VLVisionRotaryEmbedding(head_dim // 2) + + self.blocks = nn.ModuleList( + [Qwen3VLVisionBlock(config) for _ in range(config.depth)] + ) + self.merger = Qwen3VLVisionPatchMerger( + config=config, + use_postshuffle_norm=False, + ) + + self.deepstack_visual_indexes = config.deepstack_visual_indexes + self.deepstack_merger_list = nn.ModuleList( + [ + Qwen3VLVisionPatchMerger( + config=config, + use_postshuffle_norm=True, + ) + for _ in range(len(config.deepstack_visual_indexes)) + ] + ) + + self.gradient_checkpointing = False + + def rot_pos_emb(self, grid_thw: torch.Tensor) -> torch.Tensor: + merge_size = self.spatial_merge_size + + max_hw = int(grid_thw[:, 1:].max().item()) + freq_table = self.rotary_pos_emb(max_hw) # (max_hw, dim // 2) + device = freq_table.device + + total_tokens = int(torch.prod(grid_thw, dim=1).sum().item()) + pos_ids = torch.empty((total_tokens, 2), dtype=torch.long, device=device) + + offset = 0 + for num_frames, height, width in grid_thw: + merged_h, merged_w = height // merge_size, width // merge_size + + block_rows = torch.arange(merged_h, device=device) # block row indices + block_cols = torch.arange(merged_w, device=device) # block col indices + intra_row = torch.arange( + merge_size, device=device + ) # intra-block row offsets + intra_col = torch.arange( + merge_size, device=device + ) # intra-block col offsets + + # Compute full-resolution positions + row_idx = ( + block_rows[:, None, None, None] * merge_size + + intra_row[None, None, :, None] + ) + col_idx = ( + block_cols[None, :, None, None] * merge_size + + intra_col[None, None, None, :] + ) + + row_idx = row_idx.expand( + merged_h, merged_w, merge_size, merge_size + ).reshape(-1) + col_idx = col_idx.expand( + merged_h, merged_w, merge_size, merge_size + ).reshape(-1) + + coords = torch.stack((row_idx, col_idx), dim=-1) + + if num_frames > 1: + coords = coords.repeat(num_frames, 1) + + num_tokens = coords.shape[0] + pos_ids[offset : offset + num_tokens] = coords + offset += num_tokens + + embeddings = freq_table[pos_ids] # lookup rotary embeddings + embeddings = embeddings.flatten(1) + return embeddings + + def fast_pos_embed_interpolate(self, grid_thw): + grid_ts, grid_hs, grid_ws = grid_thw[:, 0], grid_thw[:, 1], grid_thw[:, 2] + + idx_list = [[] for _ in range(4)] + weight_list = [[] for _ in range(4)] + + for t, h, w in zip(grid_ts, grid_hs, grid_ws): + h_idxs = torch.linspace(0, self.num_grid_per_side - 1, h) + w_idxs = torch.linspace(0, self.num_grid_per_side - 1, w) + + h_idxs_floor = h_idxs.int() + w_idxs_floor = w_idxs.int() + h_idxs_ceil = (h_idxs.int() + 1).clip(max=self.num_grid_per_side - 1) + w_idxs_ceil = (w_idxs.int() + 1).clip(max=self.num_grid_per_side - 1) + + dh = h_idxs - h_idxs_floor + dw = w_idxs - w_idxs_floor + + base_h = h_idxs_floor * self.num_grid_per_side + base_h_ceil = h_idxs_ceil * self.num_grid_per_side + + indices = [ + (base_h[None].T + w_idxs_floor[None]).flatten(), + (base_h[None].T + w_idxs_ceil[None]).flatten(), + (base_h_ceil[None].T + w_idxs_floor[None]).flatten(), + (base_h_ceil[None].T + w_idxs_ceil[None]).flatten(), + ] + + weights = [ + ((1 - dh)[None].T * (1 - dw)[None]).flatten(), + ((1 - dh)[None].T * dw[None]).flatten(), + (dh[None].T * (1 - dw)[None]).flatten(), + (dh[None].T * dw[None]).flatten(), + ] + + for i in range(4): + idx_list[i].extend(indices[i].tolist()) + weight_list[i].extend(weights[i].tolist()) + + idx_tensor = torch.tensor( + idx_list, dtype=torch.long, device=self.pos_embed.weight.device + ) + weight_tensor = torch.tensor( + weight_list, + dtype=self.pos_embed.weight.dtype, + device=self.pos_embed.weight.device, + ) + pos_embeds = self.pos_embed(idx_tensor) * weight_tensor[:, :, None] + patch_pos_embeds = pos_embeds[0] + pos_embeds[1] + pos_embeds[2] + pos_embeds[3] + + patch_pos_embeds = patch_pos_embeds.split( + [h * w for h, w in zip(grid_hs, grid_ws)] + ) + + patch_pos_embeds_permute = [] + merge_size = self.config.spatial_merge_size + for pos_embed, t, h, w in zip(patch_pos_embeds, grid_ts, grid_hs, grid_ws): + pos_embed = pos_embed.repeat(t, 1) + pos_embed = ( + pos_embed.view( + t, h // merge_size, merge_size, w // merge_size, merge_size, -1 + ) + .permute(0, 1, 3, 2, 4, 5) + .flatten(0, 4) + ) + patch_pos_embeds_permute.append(pos_embed) + patch_pos_embeds = torch.cat(patch_pos_embeds_permute) + return patch_pos_embeds + + def forward( + self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs + ) -> torch.Tensor: + """ + Args: + hidden_states (`torch.Tensor` of shape `(seq_len, hidden_size)`): + The final hidden states of the model. + grid_thw (`torch.Tensor` of shape `(num_images_or_videos, 3)`): + The temporal, height and width of feature shape of each image in LLM. + + Returns: + `torch.Tensor`: hidden_states. + """ + hidden_states = self.patch_embed(hidden_states) + + pos_embeds = self.fast_pos_embed_interpolate(grid_thw) + hidden_states = hidden_states + pos_embeds + + rotary_pos_emb = self.rot_pos_emb(grid_thw) + + seq_len, _ = hidden_states.size() + hidden_states = hidden_states.reshape(seq_len, -1) + rotary_pos_emb = rotary_pos_emb.reshape(seq_len, -1) + emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1) + position_embeddings = (emb.cos(), emb.sin()) + + cu_seqlens = torch.repeat_interleave( + grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0] + ).cumsum( + dim=0, + # Select dtype based on the following factors: + # - FA2 requires that cu_seqlens_q must have dtype int32 + # - torch.onnx.export requires that cu_seqlens_q must have same dtype as grid_thw + # See https://github.com/huggingface/transformers/pull/34852 for more information + dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32, + ) + cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0) + + deepstack_feature_lists = [] + for layer_num, blk in enumerate(self.blocks): + hidden_states = blk( + hidden_states, + cu_seqlens=cu_seqlens, + position_embeddings=position_embeddings, + **kwargs, + ) + if layer_num in self.deepstack_visual_indexes: + deepstack_feature = self.deepstack_merger_list[ + self.deepstack_visual_indexes.index(layer_num) + ](hidden_states) + deepstack_feature_lists.append(deepstack_feature) + + hidden_states = self.merger(hidden_states) + + return hidden_states, deepstack_feature_lists + + +@auto_docstring( + custom_intro=( + "Text part of Qwen3VL, " + "not a pure text-only model, as DeepStack integrates visual features into the early hidden states." + ) +) +class Qwen3VLTextModel(Qwen3VLPreTrainedModel): + config: Qwen3VLTextConfig + _no_split_modules = ["Qwen3VLTextDecoderLayer"] + + def __init__(self, config: Qwen3VLTextConfig): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embed_tokens = nn.Embedding( + config.vocab_size, config.hidden_size, self.padding_idx + ) + self.layers = nn.ModuleList( + [ + Qwen3VLTextDecoderLayer(config, layer_idx) + for layer_idx in range(config.num_hidden_layers) + ] + ) + self.norm = Qwen3VLTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = Qwen3VLTextRotaryEmbedding(config=config) + self.gradient_checkpointing = False + + # Initialize weights and apply final processing + self.post_init() + + @merge_with_config_defaults + @auto_docstring + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Cache] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + use_cache: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + # args for deepstack + visual_pos_masks: Optional[torch.Tensor] = None, + deepstack_visual_embeds: Optional[list[torch.Tensor]] = None, + return_mid_results_layers: Optional[list] = None, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> Union[tuple, BaseModelOutputWithPast]: + r""" + cache_position (`torch.LongTensor` of shape `(seqlen,)`, *optional*): + Indices depicting the position of the input sequence tokens in the sequence. Used to index into the + key/value cache for incremental decoding. + visual_pos_masks (`torch.Tensor` of shape `(batch_size, seqlen)`, *optional*): + The mask of the visual positions. + deepstack_visual_embeds (`list[torch.Tensor]`, *optional*): + The deepstack visual embeddings. The shape is (num_layers, visual_seqlen, embed_dim). + The feature is extracted from the different visual encoder layers, and fed to the decoder + hidden states. It's from the paper DeepStack(https://arxiv.org/abs/2406.04334). + return_mid_results_layers (`list[int]`, *optional*): + Indices of decoder layers whose hidden states should be collected and returned as + intermediate results on the output's `mid_results` attribute. + """ + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError( + "You must specify exactly one of input_ids or inputs_embeds" + ) + + # torch.jit.trace() doesn't support cache objects in the output + if use_cache and past_key_values is None and not torch.jit.is_tracing(): + past_key_values = DynamicCache(config=self.config) + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + if cache_position is None: + past_seen_tokens = ( + past_key_values.get_seq_length() if past_key_values is not None else 0 + ) + cache_position = torch.arange( + past_seen_tokens, + past_seen_tokens + inputs_embeds.shape[1], + device=inputs_embeds.device, + ) + + # the hard coded `3` is for temporal, height and width. + if position_ids is None: + position_ids = cache_position.view(1, 1, -1).expand( + 3, inputs_embeds.shape[0], -1 + ) + elif position_ids.ndim == 2: + position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1) + + if position_ids.ndim == 3 and position_ids.shape[0] == 4: + text_position_ids = position_ids[0] + position_ids = position_ids[1:] + else: + text_position_ids = position_ids[0] + + attention_mask = create_causal_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + cache_position=cache_position, + past_key_values=past_key_values, + position_ids=text_position_ids, + ) + + hidden_states = inputs_embeds + mid_results = [] if return_mid_results_layers else None + + # create position embeddings to be shared across the decoder layers + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + # === Memory profiling for decoder loop (gated by DEBUG_MEM=1) === + import os as _os + + _mem_debug = _os.environ.get("DEBUG_MEM", "0") == "1" + _gc_count = 0 + + # decoder layers + for layer_idx, decoder_layer in enumerate(self.layers): + if self.gradient_checkpointing and torch.is_grad_enabled(): + # Use HuggingFace's _gradient_checkpointing_func which already has + # use_reentrant=False baked in from gradient_checkpointing_enable(). + layer_outputs = self._gradient_checkpointing_func( + decoder_layer.__call__, + hidden_states, + attention_mask=attention_mask, + position_ids=text_position_ids, + past_key_values=past_key_values, + cache_position=cache_position, + position_embeddings=position_embeddings, + ) + _gc_count += 1 + else: + layer_outputs = decoder_layer( + hidden_states, + attention_mask=attention_mask, + position_ids=text_position_ids, + past_key_values=past_key_values, + cache_position=cache_position, + position_embeddings=position_embeddings, + **kwargs, + ) + hidden_states = layer_outputs + + # Per-layer memory logging (gated by DEBUG_MEM=1) + if _mem_debug and (layer_idx % 4 == 0 or layer_idx == len(self.layers) - 1): + _a = torch.cuda.memory_allocated() / 1e9 + _rank = int(_os.environ.get("RANK", 0)) + print( + f"[MEM][rank{_rank}][decoder] layer {layer_idx:2d}/{len(self.layers)}: alloc={_a:.2f}GB", + flush=True, + ) + + # add visual features to the hidden states of first several layers + if deepstack_visual_embeds is not None and layer_idx in range( + len(deepstack_visual_embeds) + ): + hidden_states = self._deepstack_process( + hidden_states, + visual_pos_masks, + deepstack_visual_embeds[layer_idx], + ) + + if ( + return_mid_results_layers is not None + and layer_idx in return_mid_results_layers + ): + mid_results.append(hidden_states) + + _a = torch.cuda.memory_allocated() / 1e9 + if _mem_debug: + _rank = int(_os.environ.get("RANK", 0)) + print( + f"[MEM][rank{_rank}][decoder] LOOP END: alloc={_a:.2f}GB, GC_used={_gc_count}/{len(self.layers)} layers", + flush=True, + ) + + hidden_states = self.norm(hidden_states) + + output = BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + ) + output.mid_results = mid_results + return output + + def _deepstack_process( + self, + hidden_states: torch.Tensor, + visual_pos_masks: torch.Tensor, + visual_embeds: torch.Tensor, + ): + visual_pos_masks = visual_pos_masks.to(hidden_states.device) + visual_embeds = visual_embeds.to(hidden_states.device, hidden_states.dtype) + local_this = hidden_states[visual_pos_masks, :].clone() + visual_embeds + hidden_states[visual_pos_masks, :] = local_this + return hidden_states + + +class BottleneckPatchEmbed(nn.Module): + def __init__( + self, config, patch_size=16, in_chans=3, pca_dim=768, embed_dim=768, bias=True + ): + super().__init__() + self.proj1 = nn.Linear(patch_size * patch_size * in_chans, pca_dim, bias=False) + self.proj2 = nn.Linear(pca_dim, embed_dim, bias=bias) + self.initialize_weights() + + def initialize_weights(self): + w1 = self.proj1.weight.data + nn.init.xavier_uniform_(w1.view([w1.shape[0], -1])) + w2 = self.proj2.weight.data + nn.init.xavier_uniform_(w2.view([w2.shape[0], -1])) + nn.init.constant_(self.proj2.bias, 0) + + def forward(self, x): + x = self.proj2(self.proj1(x)) + return x + + +class FinalLayer(nn.Module): + def __init__(self, config, hidden_size, patch_size, out_channels): + super().__init__() + self.linear = nn.Linear( + hidden_size, patch_size * patch_size * out_channels, bias=True + ) + self.apply(self._init_weights) + + def _init_weights(self, m): + if isinstance(m, nn.Linear): + nn.init.zeros_(m.weight) + if m.bias is not None: + nn.init.constant_(m.bias, 0) + + def forward(self, x, adaln_input=None): + x = self.linear(x) + return x + + +class TimestepEmbedder(nn.Module): + """ + Embeds scalar timesteps into vector representations. + """ + + def __init__(self, config, hidden_size, frequency_embedding_size=256): + super().__init__() + self.mlp = nn.Sequential( + nn.Linear(frequency_embedding_size, hidden_size, bias=True), + nn.SiLU(), + nn.Linear(hidden_size, hidden_size, bias=True), + ) + nn.init.normal_(self.mlp[0].weight, std=0.02) + nn.init.normal_(self.mlp[2].weight, std=0.02) + self.frequency_embedding_size = frequency_embedding_size + + @staticmethod + def timestep_embedding(t, dim, max_period=10000): + """ + Create sinusoidal timestep embeddings. + :param t: a 1-D Tensor of N indices, one per batch element. + These may be fractional. + :param dim: the dimension of the output. + :param max_period: controls the minimum frequency of the embeddings. + :return: an (N, D) Tensor of positional embeddings. + """ + # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py + half = dim // 2 + freqs = torch.exp( + -math.log(max_period) + * torch.arange(start=0, end=half, dtype=torch.float32) + / half + ).to(device=t.device) + args = t[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat( + [embedding, torch.zeros_like(embedding[:, :1])], dim=-1 + ) + return embedding + + def forward(self, t): + t_freq = self.timestep_embedding(t * 1000, self.frequency_embedding_size) + t_emb = self.mlp(t_freq.to(self.mlp[0].weight.dtype)) + return t_emb + + +@auto_docstring +class Qwen3VLModel(Qwen3VLPreTrainedModel): + base_model_prefix = "" + _checkpoint_conversion_mapping = {} + # Reference: fix gemma3 grad acc #37208 + accepts_loss_kwargs = False + config: Qwen3VLConfig + _no_split_modules = ["Qwen3VLTextDecoderLayer", "Qwen3VLVisionBlock"] + + def __init__(self, config): + super().__init__(config) + self.visual = Qwen3VLVisionModel._from_config(config.vision_config) + self.language_model = Qwen3VLTextModel._from_config(config.text_config) + self.rope_deltas = None # cache rope_deltas here + + self.patch_size = 32 + self.in_channels = 3 + hidden_size = config.text_config.hidden_size + bottleneck_dim = hidden_size // 4 + + self.t_embedder1 = TimestepEmbedder(self.config, hidden_size) + self.x_embedder = BottleneckPatchEmbed( + self.config, + patch_size=self.patch_size, + in_chans=self.in_channels, + pca_dim=bottleneck_dim, + embed_dim=hidden_size, + bias=True, + ) + + # self.t_embedder2 = TimestepEmbedder(self.config, hidden_size) + self.t_embedder2 = None + self.final_layer2 = FinalLayer( + self.config, + hidden_size=hidden_size, + patch_size=self.patch_size, + out_channels=self.in_channels, + ) + self.tms_token_id = 151673 + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.language_model.get_input_embeddings() + + def set_input_embeddings(self, value): + self.language_model.set_input_embeddings(value) + + def set_decoder(self, decoder): + self.language_model = decoder + + def get_decoder(self): + return self.language_model + + def get_rope_index( + self, + input_ids: Optional[torch.LongTensor] = None, + image_grid_thw: Optional[torch.LongTensor] = None, + video_grid_thw: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Different from the original implementation, Qwen3VL use timestamps rather than absolute time position ids.""" + + # Since we use timestamps to seperate videos, like , the video_grid_thw should also be split + if video_grid_thw is not None: + video_grid_thw = torch.repeat_interleave( + video_grid_thw, video_grid_thw[:, 0], dim=0 + ) + video_grid_thw[:, 0] = 1 + + spatial_merge_size = self.config.vision_config.spatial_merge_size + image_token_id = self.config.image_token_id + video_token_id = self.config.video_token_id + vision_start_token_id = self.config.vision_start_token_id + mrope_position_deltas = [] + if input_ids is not None and ( + image_grid_thw is not None or video_grid_thw is not None + ): + total_input_ids = input_ids + if attention_mask is None: + attention_mask = torch.ones_like(total_input_ids) + position_ids = torch.ones( + 3, + input_ids.shape[0], + input_ids.shape[1], + dtype=input_ids.dtype, + device=input_ids.device, + ) + image_index, video_index = 0, 0 + attention_mask = attention_mask.to(total_input_ids.device) + for i, input_ids in enumerate(total_input_ids): + input_ids = input_ids[attention_mask[i] == 1] + image_nums, video_nums = 0, 0 + vision_start_indices = torch.argwhere( + input_ids == vision_start_token_id + ).squeeze(1) + vision_tokens = input_ids[vision_start_indices + 1] + image_nums = (vision_tokens == image_token_id).sum() + video_nums = (vision_tokens == video_token_id).sum() + input_tokens = input_ids.tolist() + llm_pos_ids_list: list = [] + st = 0 + remain_images, remain_videos = image_nums, video_nums + for _ in range(image_nums + video_nums): + if image_token_id in input_tokens and remain_images > 0: + ed_image = input_tokens.index(image_token_id, st) + else: + ed_image = len(input_tokens) + 1 + if video_token_id in input_tokens and remain_videos > 0: + ed_video = input_tokens.index(video_token_id, st) + else: + ed_video = len(input_tokens) + 1 + if ed_image < ed_video: + t, h, w = ( + image_grid_thw[image_index][0], + image_grid_thw[image_index][1], + image_grid_thw[image_index][2], + ) + image_index += 1 + remain_images -= 1 + ed = ed_image + + else: + t, h, w = ( + video_grid_thw[video_index][0], + video_grid_thw[video_index][1], + video_grid_thw[video_index][2], + ) + video_index += 1 + remain_videos -= 1 + ed = ed_video + llm_grid_t, llm_grid_h, llm_grid_w = ( + t.item(), + h.item() // spatial_merge_size, + w.item() // spatial_merge_size, + ) + text_len = ed - st + + st_idx = ( + llm_pos_ids_list[-1].max() + 1 + if len(llm_pos_ids_list) > 0 + else 0 + ) + llm_pos_ids_list.append( + torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx + ) + + # t_index is always 0 because llm_grid_t is always 1 (we use timestamps to encode the temporal information for videos) + t_index = ( + torch.arange(llm_grid_t) + .view(-1, 1) + .expand(-1, llm_grid_h * llm_grid_w) + .flatten() + ) + h_index = ( + torch.arange(llm_grid_h) + .view(1, -1, 1) + .expand(llm_grid_t, -1, llm_grid_w) + .flatten() + ) + w_index = ( + torch.arange(llm_grid_w) + .view(1, 1, -1) + .expand(llm_grid_t, llm_grid_h, -1) + .flatten() + ) + llm_pos_ids_list.append( + torch.stack([t_index, h_index, w_index]) + text_len + st_idx + ) + st = ed + llm_grid_t * llm_grid_h * llm_grid_w + + if st < len(input_tokens): + st_idx = ( + llm_pos_ids_list[-1].max() + 1 + if len(llm_pos_ids_list) > 0 + else 0 + ) + text_len = len(input_tokens) - st + llm_pos_ids_list.append( + torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx + ) + + llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) + position_ids[..., i, attention_mask[i] == 1] = llm_positions.to( + position_ids.device + ) + mrope_position_deltas.append( + llm_positions.max() + 1 - len(total_input_ids[i]) + ) + mrope_position_deltas = torch.tensor( + mrope_position_deltas, device=input_ids.device + ).unsqueeze(1) + return position_ids, mrope_position_deltas + else: + if attention_mask is not None: + position_ids = attention_mask.long().cumsum(-1) - 1 + position_ids.masked_fill_(attention_mask == 0, 1) + position_ids = ( + position_ids.unsqueeze(0) + .expand(3, -1, -1) + .to(attention_mask.device) + ) + max_position_ids = position_ids.max(0, keepdim=False)[0].max( + -1, keepdim=True + )[0] + mrope_position_deltas = max_position_ids + 1 - attention_mask.shape[-1] + else: + position_ids = ( + torch.arange(input_ids.shape[1], device=input_ids.device) + .view(1, 1, -1) + .expand(3, input_ids.shape[0], -1) + ) + mrope_position_deltas = torch.zeros( + [input_ids.shape[0], 1], + device=input_ids.device, + dtype=input_ids.dtype, + ) + + return position_ids, mrope_position_deltas + + def get_video_features( + self, + pixel_values_videos: torch.FloatTensor, + video_grid_thw: Optional[torch.LongTensor] = None, + ): + """ + Encodes videos into continuous embeddings that can be forwarded to the language model. The deepstack visual features are also returned. + + Args: + pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input videos. + video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): + The temporal, height and width of feature shape of each video in LLM. + """ + # Same implementation as for images + return self.get_image_features(pixel_values_videos, video_grid_thw) + + def get_image_features( + self, + pixel_values: torch.FloatTensor, + image_grid_thw: Optional[torch.LongTensor] = None, + ): + """ + Encodes images into continuous embeddings that can be forwarded to the language model. The deepstack visual features are also returned. + + Args: + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input images. + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. + """ + pixel_values = pixel_values.type(self.visual.dtype) + image_embeds, deepstack_image_embeds = self.visual( + pixel_values, grid_thw=image_grid_thw + ) + split_sizes = ( + image_grid_thw.prod(-1) // self.visual.spatial_merge_size**2 + ).tolist() + image_embeds = torch.split(image_embeds, split_sizes) + return image_embeds, deepstack_image_embeds + + def get_placeholder_mask( + self, + input_ids: torch.LongTensor, + inputs_embeds: torch.FloatTensor, + image_features: Optional[torch.FloatTensor] = None, + video_features: Optional[torch.FloatTensor] = None, + ): + """ + Obtains multimodal placeholder mask from `input_ids` or `inputs_embeds`, and checks that the placeholder token count is + equal to the length of multimodal features. If the lengths are different, an error is raised. + """ + if input_ids is None: + special_image_mask = inputs_embeds == self.get_input_embeddings()( + torch.tensor( + self.config.image_token_id, + dtype=torch.long, + device=inputs_embeds.device, + ) + ) + special_image_mask = special_image_mask.all(-1) + special_video_mask = inputs_embeds == self.get_input_embeddings()( + torch.tensor( + self.config.video_token_id, + dtype=torch.long, + device=inputs_embeds.device, + ) + ) + special_video_mask = special_video_mask.all(-1) + else: + special_image_mask = input_ids == self.config.image_token_id + special_video_mask = input_ids == self.config.video_token_id + + n_image_tokens = special_image_mask.sum() + special_image_mask = ( + special_image_mask.unsqueeze(-1) + .expand_as(inputs_embeds) + .to(inputs_embeds.device) + ) + if ( + image_features is not None + and inputs_embeds[special_image_mask].numel() != image_features.numel() + ): + raise ValueError( + f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {image_features.shape[0]}" + ) + + n_video_tokens = special_video_mask.sum() + special_video_mask = ( + special_video_mask.unsqueeze(-1) + .expand_as(inputs_embeds) + .to(inputs_embeds.device) + ) + if ( + video_features is not None + and inputs_embeds[special_video_mask].numel() != video_features.numel() + ): + raise ValueError( + f"Videos features and video tokens do not match: tokens: {n_video_tokens}, features {video_features.shape[0]}" + ) + + return special_image_mask, special_video_mask + + def _run_decoder_flash( + self, inputs_embeds, position_ids, token_types, return_mid_results_layers=None + ): + """Run decoder layers with two-pass attention. + + Replicates the Megatron attention pattern: + 1. Causal attention on AR tokens only (text) + 2. Full (bidirectional) attention on ALL tokens + 3. Replace AR positions with causal result (index_copy) + + This ensures AR tokens only attend causally to other AR tokens, + while gen tokens attend bidirectionally to everything. + + Uses the transformers attention dispatch (ALL_ATTENTION_FUNCTIONS), + so any backend works (sdpa by default, flash_attention_2 if + activated via config._attn_implementation). + + Args: + inputs_embeds: [batch, total_seq_len, hidden] + position_ids: [3, batch, total_seq_len] - 3D RoPE positions + token_types: [batch, total_seq_len] - 0=AR, 1=gen + """ + text_model = self.language_model + + # Compute rotary position embeddings + if position_ids.ndim == 2: + position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1) + elif position_ids.ndim == 3 and position_ids.shape[0] == 4: + position_ids = position_ids[1:] # drop text_position_ids dim + position_embeddings = text_model.rotary_emb(inputs_embeds, position_ids) + cos, sin = position_embeddings + + # Precompute AR token indices (same layout for all batch items) + is_gen = token_types[0].bool() # [seq_len] + idx_ar = torch.nonzero(~is_gen, as_tuple=False).squeeze(-1) # [n_ar] + + hidden_states = inputs_embeds + mid_results = [] if return_mid_results_layers else None + + use_gc = text_model.gradient_checkpointing and torch.is_grad_enabled() + + def _two_pass_layer_forward(hidden_states, decoder_layer, cos, sin, idx_ar): + """Two-pass attention layer forward compatible with FSDP2. + + Calls decoder_layer(...) through its __call__ to trigger FSDP hooks + (which swap DTensor parameters to plain tensors), with self_attn.forward + temporarily replaced by a custom two-pass attention implementation + that goes through the transformers attention dispatch. + """ + original_attn_forward = decoder_layer.self_attn.forward + + def _custom_two_pass_attn( + hidden_states, position_embeddings, attention_mask=None, **kwargs + ): + attn = decoder_layer.self_attn + input_shape = hidden_states.shape[:-1] + head_dim = attn.head_dim + hidden_shape = (*input_shape, -1, head_dim) + + # Q, K, V projections in [B, H, S, D] + q = attn.q_norm( + attn.q_proj(hidden_states).view(hidden_shape) + ).transpose(1, 2) + k = attn.k_norm( + attn.k_proj(hidden_states).view(hidden_shape) + ).transpose(1, 2) + v = attn.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) + + # Apply rotary position embedding + cos_pe, sin_pe = position_embeddings + q, k = apply_rotary_pos_emb(q, k, cos_pe, sin_pe) + + scaling = head_dim**-0.5 + + # Attention dispatch — sdpa/eager by default, flash_attention_2 + # if activated via config._attn_implementation. + attention_interface: Callable = eager_attention_forward + if attn.config._attn_implementation != "eager": + attention_interface = ALL_ATTENTION_FUNCTIONS[ + attn.config._attn_implementation + ] + + # --- Two-pass attention --- + # Pass 1: causal on AR tokens only (slice on seq dim) + q_ar = q[:, :, idx_ar].contiguous() + k_ar = k[:, :, idx_ar].contiguous() + v_ar = v[:, :, idx_ar].contiguous() + out_ar, _ = attention_interface( + attn, + q_ar, + k_ar, + v_ar, + attention_mask=None, + dropout=0.0, + scaling=scaling, + is_causal=True, + ) # [B, n_ar, H, D] + + # Pass 2: full (bidirectional) on all tokens + out_full, _ = attention_interface( + attn, + q, + k, + v, + attention_mask=None, + dropout=0.0, + scaling=scaling, + is_causal=False, + ) # [B, S, H, D] + + # Replace AR positions with causal result + out_full = out_full.clone() + out_full[:, idx_ar] = out_ar + + # Output projection + attn_output = out_full.reshape(*input_shape, -1).contiguous() + attn_output = attn.o_proj(attn_output) + return attn_output, None + + # Temporarily disable gradient checkpointing on the decoder layer + # to avoid nested checkpointing (the outer loop handles GC). + _saved_gc = decoder_layer.gradient_checkpointing + decoder_layer.gradient_checkpointing = False + decoder_layer.self_attn.forward = _custom_two_pass_attn + try: + hidden_states = decoder_layer( + hidden_states, + position_embeddings=(cos, sin), + ) + finally: + decoder_layer.self_attn.forward = original_attn_forward + decoder_layer.gradient_checkpointing = _saved_gc + + return hidden_states + + for layer_idx, decoder_layer in enumerate(text_model.layers): + if use_gc: + hidden_states = torch.utils.checkpoint.checkpoint( + _two_pass_layer_forward, + hidden_states, + decoder_layer, + cos, + sin, + idx_ar, + use_reentrant=False, + ) + else: + hidden_states = _two_pass_layer_forward( + hidden_states, + decoder_layer, + cos, + sin, + idx_ar, + ) + + if ( + return_mid_results_layers is not None + and layer_idx in return_mid_results_layers + ): + mid_results.append(hidden_states) + + # Final layer norm + hidden_states = text_model.norm(hidden_states) + return hidden_states, mid_results + + def _forward_generation( + self, + input_ids, + position_ids, + vinputs, + timestep, + token_types, + attention_mask=None, + pixel_values=None, + pixel_values_videos=None, + image_grid_thw=None, + video_grid_thw=None, + use_flash_attn=False, + return_mid_results_layers=None, + **kwargs, + ): + """Forward pass for image generation (denoising step). + + Args: + input_ids: [batch, txt_seq_len] - text token IDs (without image tokens) + position_ids: [3, batch, total_seq_len] - 3D RoPE positions covering text+image + vinputs: [batch, img_tokens, patch_dim] - patchified noise input + timestep: [batch] - timestep values (scalar per sample) + token_types: [batch, total_seq_len] or [total_seq_len, 1] - 0=AR, >0=gen + attention_mask: ignored (created internally for non-flash path) + pixel_values: optional image pixel values for conditioned generation + pixel_values_videos: optional video pixel values + image_grid_thw: optional image grid info + video_grid_thw: optional video grid info + use_flash_attn: if True, use flash attention with two-pass approach + + Returns: + Qwen3VLModelOutputWithPast with x_pred field set. + """ + # 1. Get text token embeddings + inputs_embeds = self.get_input_embeddings()( + input_ids + ) # [batch, txt_seq_len, hidden] + + # 2. Process image/video embeddings if present (for image-conditioned generation) + if pixel_values is not None: + image_embeds, _ = self.get_image_features(pixel_values, image_grid_thw) + image_embeds = torch.cat(image_embeds, dim=0).to( + inputs_embeds.device, inputs_embeds.dtype + ) + image_mask, _ = self.get_placeholder_mask( + input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds + ) + inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) + elif torch.is_grad_enabled(): + # t2i task: no pixel_values, but we must run the vision encoder with a + # tiny dummy input so that EVERY rank has non-None (zero) gradients for + # vision-encoder parameters. This keeps the FSDP reduce-scatter and the + # replicate-group all-reduce symmetric across t2i and ref-task ranks, + # preventing collective hangs at backward / clip_grad_norm_. + # The dummy output is zeroed out before being added to inputs_embeds, so + # the forward result is numerically identical to the no-pixel_values path. + pe = self.visual.patch_embed # PatchEmbed + t_sz = pe.temporal_patch_size # e.g. 2 + m_sz = self.visual.spatial_merge_size # e.g. 2 + n_patches = t_sz * m_sz * m_sz + patch_dim = pe.in_channels * t_sz * pe.patch_size * pe.patch_size + fake_pv = torch.zeros( + n_patches, + patch_dim, + device=inputs_embeds.device, + dtype=pe.proj.weight.dtype, + ) + fake_grid = torch.tensor( + [[t_sz, m_sz, m_sz]], dtype=torch.long, device=inputs_embeds.device + ) + fake_embs, _ = self.get_image_features(fake_pv, fake_grid) + fake_embs = torch.cat(fake_embs, dim=0).to(inputs_embeds.dtype) + # Multiply by a zero tensor (same dtype/device) so the gradient path + # through the vision encoder is live but the numerical contribution is 0. + inputs_embeds = inputs_embeds + fake_embs.sum() * inputs_embeds.new_zeros( + [] + ) + + if pixel_values_videos is not None: + video_embeds, _ = self.get_video_features( + pixel_values_videos, video_grid_thw + ) + video_embeds = torch.cat(video_embeds, dim=0).to( + inputs_embeds.device, inputs_embeds.dtype + ) + _, video_mask = self.get_placeholder_mask( + input_ids, inputs_embeds=inputs_embeds, video_features=video_embeds + ) + inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds) + + # 3. Embed timestep and replace tms_token positions + if isinstance(timestep, list): + timestep = torch.cat(timestep, dim=0) + timestep = timestep.to(inputs_embeds.device) + t_emb = self.t_embedder1(timestep) # [batch, hidden] + + tms_mask = input_ids == self.tms_token_id # [batch, txt_seq_len] + tms_mask_3d = tms_mask.unsqueeze(-1).expand_as(inputs_embeds) + t_emb_expanded = t_emb.unsqueeze(1).expand_as(inputs_embeds) + inputs_embeds = torch.where(tms_mask_3d, t_emb_expanded, inputs_embeds) + + # 4. Embed vinputs and append to sequence + if isinstance(vinputs, list): + vinputs = torch.cat(vinputs, dim=0) + vinputs = vinputs.to(inputs_embeds.device) + vinputs_embedded = self.x_embedder(vinputs).to( + inputs_embeds.dtype + ) # [batch, img_tokens, hidden] + inputs_embeds = torch.cat( + [inputs_embeds, vinputs_embedded], dim=1 + ) # [batch, total_seq_len, hidden] + + batch_size, total_seq_len, _ = inputs_embeds.shape + + # 5. Parse token_types to [batch, total_seq_len] + if isinstance(token_types, list): + token_types = torch.cat(token_types, dim=0) + token_types = token_types.to(inputs_embeds.device) + if token_types.dim() == 1: + token_types = token_types.unsqueeze(0) + elif ( + token_types.dim() == 2 + and token_types.shape[-1] == 1 + and token_types.shape[0] == total_seq_len + ): + # [total_seq_len, 1] -> [1, total_seq_len] + token_types = token_types.squeeze(-1).unsqueeze(0) + if token_types.shape[0] == 1 and batch_size > 1: + token_types = token_types.expand(batch_size, -1) + + # 6. Forward through decoder + mid_results = None + import os as _os2 + + _mem_debug2 = _os2.environ.get("DEBUG_MEM", "0") == "1" + + if _mem_debug2: + _rank2 = int(_os2.environ.get("RANK", 0)) + _a = torch.cuda.memory_allocated() / 1e9 + print( + f"[MEM][rank{_rank2}][_forward_gen] before decoder: alloc={_a:.2f}GB, " + f"total_seq_len={total_seq_len}, batch={batch_size}, flash={use_flash_attn}", + flush=True, + ) + + if use_flash_attn: + # Flash attention: two-pass approach (causal on AR + full on all → index_copy) + hidden_states, mid_results = self._run_decoder_flash( + inputs_embeds, + position_ids, + token_types, + return_mid_results_layers=return_mid_results_layers, + ) + else: + # Standard path: 4D attention mask (causal for AR, full for gen tokens) + dtype = inputs_embeds.dtype + min_val = torch.finfo(dtype).min + attn_masks = [] + for b in range(batch_size): + causal = torch.full( + (total_seq_len, total_seq_len), + min_val, + device=inputs_embeds.device, + dtype=dtype, + ) + causal = torch.triu( + causal, diagonal=1 + ) # lower tri + diag = 0 (allowed) + gen_positions = token_types[b].bool() # [total_seq_len] + causal[gen_positions, :] = 0 # gen tokens attend to everything + attn_masks.append(causal) + attention_mask_4d = torch.stack(attn_masks, dim=0).unsqueeze( + 1 + ) # [batch, 1, seq, seq] + + if _mem_debug2: + _rank2 = int(_os2.environ.get("RANK", 0)) + _mask_gb = ( + attention_mask_4d.element_size() * attention_mask_4d.numel() / 1e9 + ) + _a = torch.cuda.memory_allocated() / 1e9 + print( + f"[MEM][rank{_rank2}][_forward_gen] attn_mask_4d: shape={list(attention_mask_4d.shape)}, " + f"size={_mask_gb:.3f}GB, alloc={_a:.2f}GB", + flush=True, + ) + + outputs = self.language_model( + input_ids=None, + position_ids=position_ids, + attention_mask=attention_mask_4d, + inputs_embeds=inputs_embeds, + use_cache=False, + return_mid_results_layers=return_mid_results_layers, + ) + hidden_states = outputs.last_hidden_state + if hasattr(outputs, "mid_results"): + mid_results = outputs.mid_results + + # 7. Apply final layer to get pixel predictions + x_pred = self.final_layer2(hidden_states) # [batch, total_seq_len, out_dim] + + return Qwen3VLModelOutputWithPast( + last_hidden_state=hidden_states, + x_pred=x_pred, + mid_results=mid_results, + ) + + @auto_docstring + @merge_with_config_defaults + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Cache] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + pixel_values: Optional[torch.Tensor] = None, + pixel_values_videos: Optional[torch.FloatTensor] = None, + image_grid_thw: Optional[torch.LongTensor] = None, + video_grid_thw: Optional[torch.LongTensor] = None, + cache_position: Optional[torch.LongTensor] = None, + vinputs: Optional[torch.Tensor] = None, + timestep: Optional[torch.Tensor] = None, + token_types: Optional[torch.Tensor] = None, + use_flash_attn: bool = False, + return_mid_results_layers: Optional[list] = None, + **kwargs: Unpack[TransformersKwargs], + ) -> Union[tuple, Qwen3VLModelOutputWithPast]: + r""" + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. + video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): + The temporal, height and width of feature shape of each video in LLM. + cache_position (`torch.LongTensor` of shape `(seqlen,)`, *optional*): + Indices depicting the position of the input sequence tokens in the sequence. Used to index into the + key/value cache for incremental decoding. + vinputs (`torch.Tensor`, *optional*): + Visual inputs for the generation pathway. When provided, the forward call is dispatched to + `_forward_generation` instead of the standard understanding path. + timestep (`torch.Tensor`, *optional*): + Diffusion timestep tensor passed through to the generation forward when `vinputs` is provided. + token_types (`torch.Tensor`, *optional*): + Per-token type identifiers used by the generation forward to distinguish text vs. visual tokens. + use_flash_attn (`bool`, *optional*, defaults to `False`): + Whether to use a flash-attention kernel in the generation forward path. + return_mid_results_layers (`list[int]`, *optional*): + Indices of decoder layers whose hidden states should be collected and returned as + intermediate results on the output's `mid_results` attribute. + """ + # Dispatch to generation forward if vinputs is provided + if vinputs is not None: + return self._forward_generation( + input_ids=input_ids, + position_ids=position_ids, + vinputs=vinputs, + timestep=timestep, + token_types=token_types, + attention_mask=attention_mask, + pixel_values=pixel_values, + pixel_values_videos=pixel_values_videos, + image_grid_thw=image_grid_thw, + video_grid_thw=video_grid_thw, + use_flash_attn=use_flash_attn, + return_mid_results_layers=return_mid_results_layers, + **kwargs, + ) + + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError( + "You must specify exactly one of input_ids or inputs_embeds" + ) + + if inputs_embeds is None: + inputs_embeds = self.get_input_embeddings()(input_ids) + + image_mask = None + video_mask = None + + if pixel_values is not None: + image_embeds, deepstack_image_embeds = self.get_image_features( + pixel_values, image_grid_thw + ) + image_embeds = torch.cat(image_embeds, dim=0).to( + inputs_embeds.device, inputs_embeds.dtype + ) + image_mask, _ = self.get_placeholder_mask( + input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds + ) + inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) + + if pixel_values_videos is not None: + video_embeds, deepstack_video_embeds = self.get_video_features( + pixel_values_videos, video_grid_thw + ) + video_embeds = torch.cat(video_embeds, dim=0).to( + inputs_embeds.device, inputs_embeds.dtype + ) + _, video_mask = self.get_placeholder_mask( + input_ids, inputs_embeds=inputs_embeds, video_features=video_embeds + ) + inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds) + + visual_pos_masks = None + deepstack_visual_embeds = None + if image_mask is not None and video_mask is not None: + # aggregate visual_pos_masks and deepstack_visual_embeds + image_mask = image_mask[..., 0] + video_mask = video_mask[..., 0] + visual_pos_masks = image_mask | video_mask + deepstack_visual_embeds = [] + image_mask_joint = image_mask[visual_pos_masks] + video_mask_joint = video_mask[visual_pos_masks] + for img_embed, vid_embed in zip( + deepstack_image_embeds, deepstack_video_embeds + ): + embed_joint = img_embed.new_zeros( + visual_pos_masks.sum(), img_embed.shape[-1] + ).to(img_embed.device) + embed_joint[image_mask_joint, :] = img_embed + embed_joint[video_mask_joint, :] = vid_embed + deepstack_visual_embeds.append(embed_joint) + elif image_mask is not None: + image_mask = image_mask[..., 0] + visual_pos_masks = image_mask + deepstack_visual_embeds = deepstack_image_embeds + elif video_mask is not None: + video_mask = video_mask[..., 0] + visual_pos_masks = video_mask + deepstack_visual_embeds = deepstack_video_embeds + + if position_ids is None: + attention_mask_tensor = ( + attention_mask + if not isinstance(attention_mask, dict) + else attention_mask["full_attention"] + ) + if attention_mask_tensor is not None and attention_mask_tensor.ndim == 4: + attention_mask_tensor = torch.diagonal( + attention_mask_tensor[:, 0], dim1=1, dim2=2 + ) + # Only apply conversion for floating point tensors (inverted masks) + if attention_mask_tensor.dtype.is_floating_point: + attention_mask_tensor = ( + attention_mask_tensor + / torch.finfo(attention_mask_tensor.dtype).min + ) + attention_mask_tensor = (1.0 - attention_mask_tensor).int() + + # Calculate RoPE index once per generation in the pre-fill stage only. + # When compiling, we can't check tensor values thus we check only input length + # It is safe to assume that `length!=1` means we're in pre-fill because compiled + # models currently cannot do asssisted decoding + prefill_compiled_stage = is_torchdynamo_compiling() and ( + (input_ids is not None and input_ids.shape[1] != 1) + or (inputs_embeds is not None and inputs_embeds.shape[1] != 1) + ) + prefill_noncompiled_stage = not is_torchdynamo_compiling() and ( + (cache_position is not None and cache_position[0] == 0) + or (past_key_values is None or past_key_values.get_seq_length() == 0) + ) + if ( + prefill_compiled_stage or prefill_noncompiled_stage + ) or self.rope_deltas is None: + position_ids, rope_deltas = self.get_rope_index( + input_ids, + image_grid_thw, + video_grid_thw, + attention_mask=attention_mask_tensor, + ) + self.rope_deltas = rope_deltas + # then use the prev pre-calculated rope-deltas to get the correct position ids + else: + batch_size, seq_length, _ = inputs_embeds.shape + delta = ( + (cache_position[0] + self.rope_deltas).to(inputs_embeds.device) + if cache_position is not None + else 0 + ) + position_ids = torch.arange(seq_length, device=inputs_embeds.device) + position_ids = position_ids.view(1, -1).expand(batch_size, -1) + if cache_position is not None: # otherwise `deltas` is an int `0` + delta = delta.repeat_interleave(batch_size // delta.shape[0], dim=0) + position_ids = position_ids.add(delta) + position_ids = position_ids.unsqueeze(0).expand(3, -1, -1) + + outputs = self.language_model( + input_ids=None, + position_ids=position_ids, + attention_mask=attention_mask, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + cache_position=cache_position, + visual_pos_masks=visual_pos_masks, + deepstack_visual_embeds=deepstack_visual_embeds, + **kwargs, + ) + + return Qwen3VLModelOutputWithPast( + last_hidden_state=outputs.last_hidden_state, + past_key_values=outputs.past_key_values, + rope_deltas=self.rope_deltas, + ) + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for Qwen3VL causal language model (or autoregressive) outputs. + """ +) +class Qwen3VLCausalLMOutputWithPast(ModelOutput): + r""" + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): + Language modeling loss (for next-token prediction). + logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`): + Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax). + past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache). + + Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see + `past_key_values` input) to speed up sequential decoding. + rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*): + The rope index difference between sequence length and multimodal rope. + """ + + loss: Optional[torch.FloatTensor] = None + logits: Optional[torch.FloatTensor] = None + past_key_values: Optional[Cache] = None + hidden_states: Optional[tuple[torch.FloatTensor]] = None + attentions: Optional[tuple[torch.FloatTensor]] = None + rope_deltas: Optional[torch.LongTensor] = None + x_pred: Optional[torch.FloatTensor] = None + mid_results: Optional[list] = None + + +class Qwen3VLForConditionalGeneration(Qwen3VLPreTrainedModel, GenerationMixin): + _checkpoint_conversion_mapping = {} + _tied_weights_keys = {"lm_head.weight": "model.language_model.embed_tokens.weight"} + # Reference: fix gemma3 grad acc #37208 + accepts_loss_kwargs = False + config: Qwen3VLConfig + + def __init__(self, config): + super().__init__(config) + self.model = Qwen3VLModel(config) + self.lm_head = nn.Linear( + config.text_config.hidden_size, config.text_config.vocab_size, bias=False + ) + + self.post_init() + + def get_input_embeddings(self): + return self.model.get_input_embeddings() + + def set_input_embeddings(self, value): + self.model.set_input_embeddings(value) + + def set_decoder(self, decoder): + self.model.set_decoder(decoder) + + def get_decoder(self): + return self.model.get_decoder() + + def get_video_features( + self, + pixel_values_videos: torch.FloatTensor, + video_grid_thw: Optional[torch.LongTensor] = None, + ): + return self.model.get_video_features(pixel_values_videos, video_grid_thw) + + def get_image_features( + self, + pixel_values: torch.FloatTensor, + image_grid_thw: Optional[torch.LongTensor] = None, + ): + return self.model.get_image_features(pixel_values, image_grid_thw) + + # Make modules available through conditional class for BC + @property + def language_model(self): + return self.model.language_model + + @property + def visual(self): + return self.model.visual + + def enable_gradient_checkpointing(self): + self.model.visual.gradient_checkpointing_enable() + self.model.language_model.gradient_checkpointing_enable() + + @merge_with_config_defaults + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Cache] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + pixel_values: Optional[torch.Tensor] = None, + pixel_values_videos: Optional[torch.FloatTensor] = None, + image_grid_thw: Optional[torch.LongTensor] = None, + video_grid_thw: Optional[torch.LongTensor] = None, + cache_position: Optional[torch.LongTensor] = None, + logits_to_keep: Union[int, torch.Tensor] = 0, + vinputs: Optional[torch.Tensor] = None, + timestep: Optional[torch.Tensor] = None, + token_types: Optional[torch.Tensor] = None, + use_flash_attn: bool = False, + return_mid_results_layers: Optional[list] = None, + **kwargs: Unpack[TransformersKwargs], + ) -> Union[tuple, Qwen3VLCausalLMOutputWithPast]: + r""" + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. + video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): + The temporal, height and width of feature shape of each video in LLM. + + Example: + TODO: Add example + """ + outputs = self.model( + input_ids=input_ids, + pixel_values=pixel_values, + pixel_values_videos=pixel_values_videos, + image_grid_thw=image_grid_thw, + video_grid_thw=video_grid_thw, + position_ids=position_ids, + attention_mask=attention_mask, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + cache_position=cache_position, + vinputs=vinputs, + timestep=timestep, + token_types=token_types, + use_flash_attn=use_flash_attn, + return_mid_results_layers=return_mid_results_layers, + **kwargs, + ) + + # Generation path: return x_pred directly + if vinputs is not None: + return Qwen3VLCausalLMOutputWithPast( + x_pred=outputs.x_pred, + mid_results=outputs.mid_results + if hasattr(outputs, "mid_results") + else None, + ) + + hidden_states = outputs[0] + + # Only compute necessary logits, and do not upcast them to float if we are not computing the loss + slice_indices = ( + slice(-logits_to_keep, None) + if isinstance(logits_to_keep, int) + else logits_to_keep + ) + logits = self.lm_head(hidden_states[:, slice_indices, :]) + + loss = None + if labels is not None: + loss = self.loss_function( + logits=logits, + labels=labels, + vocab_size=self.config.text_config.vocab_size, + ) + + return Qwen3VLCausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + rope_deltas=outputs.rope_deltas, + ) + + def prepare_inputs_for_generation( + self, + input_ids, + past_key_values=None, + attention_mask=None, + inputs_embeds=None, + cache_position=None, + position_ids=None, + use_cache=True, + pixel_values=None, + pixel_values_videos=None, + image_grid_thw=None, + video_grid_thw=None, + **kwargs, + ): + # Overwritten -- in specific circumstances we don't want to forward image inputs to the model + + model_inputs = super().prepare_inputs_for_generation( + input_ids, + past_key_values=past_key_values, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + cache_position=cache_position, + position_ids=position_ids, + pixel_values=pixel_values, + pixel_values_videos=pixel_values_videos, + image_grid_thw=image_grid_thw, + video_grid_thw=video_grid_thw, + use_cache=use_cache, + **kwargs, + ) + + # Qwen3VL position_ids are prepareed with rope_deltas in forward + model_inputs["position_ids"] = None + + if cache_position[0] != 0: + model_inputs["pixel_values"] = None + model_inputs["pixel_values_videos"] = None + + return model_inputs + + def _get_image_nums_and_video_nums( + self, + input_ids: Optional[torch.LongTensor], + inputs_embeds: Optional[torch.Tensor] = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """ + Get the number of images and videos for each sample to calculate the separation length of the sample tensor. + These parameters are not passed through the processor to avoid unpredictable impacts from interface modifications. + + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. + + Returns: + image_nums (`torch.LongTensor` of shape `(batch_size, num_images_sample)`) + video_nums (`torch.LongTensor` of shape `(batch_size, num_videos_sample)`) + """ + image_token_id = self.config.image_token_id + video_token_id = self.config.video_token_id + vision_start_token_id = self.config.vision_start_token_id + + if inputs_embeds is not None: + vision_start_mask = ( + inputs_embeds + == self.get_input_embeddings()( + torch.tensor( + vision_start_token_id, + dtype=torch.long, + device=inputs_embeds.device, + ) + ) + )[..., 0] + image_mask = ( + inputs_embeds + == self.get_input_embeddings()( + torch.tensor( + image_token_id, dtype=torch.long, device=inputs_embeds.device + ) + ) + )[..., 0] + video_mask = ( + inputs_embeds + == self.get_input_embeddings()( + torch.tensor( + video_token_id, dtype=torch.long, device=inputs_embeds.device + ) + ) + )[..., 0] + else: + vision_start_mask = input_ids == vision_start_token_id + image_mask = input_ids == image_token_id + video_mask = input_ids == video_token_id + + vision_first_mask = torch.roll(vision_start_mask, shifts=1, dims=1) + image_nums = torch.sum(vision_first_mask & image_mask, dim=1) + video_nums = torch.sum(vision_first_mask & video_mask, dim=1) + + return image_nums, video_nums + + def _expand_inputs_for_generation( + self, + expand_size: int = 1, + is_encoder_decoder: bool = False, + input_ids: Optional[torch.LongTensor] = None, + **model_kwargs, + ) -> tuple[torch.LongTensor, dict[str, Any]]: + # Overwritten -- Support for expanding tensors without a batch size dimension + # e.g., pixel_values, image_grid_thw, pixel_values_videos, video_grid_thw, second_per_grid_t + # pixel_values.shape[0] is sum(seqlen_images for samples) + # image_grid_thw.shape[0] is sum(num_images for samples) + + if expand_size == 1: + return input_ids, model_kwargs + + visual_keys = [ + "pixel_values", + "image_grid_thw", + "pixel_values_videos", + "video_grid_thw", + "second_per_grid_ts", + ] + + def _expand_dict_for_generation_visual(dict_to_expand): + image_grid_thw = model_kwargs.get("image_grid_thw", None) + video_grid_thw = model_kwargs.get("video_grid_thw", None) + image_nums, video_nums = self._get_image_nums_and_video_nums( + input_ids, inputs_embeds=model_kwargs.get("inputs_embeds", None) + ) + + def _repeat_interleave_samples(x, lengths, repeat_times): + samples = torch.split(x, lengths) + repeat_args = [repeat_times] + [1] * (x.dim() - 1) + result = torch.cat( + [sample.repeat(*repeat_args) for sample in samples], dim=0 + ) + return result + + for key in dict_to_expand: + if key == "pixel_values": + # split images into samples + samples = torch.split(image_grid_thw, list(image_nums)) + # compute the sequence length of images for each sample + lengths = [torch.prod(sample, dim=1).sum() for sample in samples] + dict_to_expand[key] = _repeat_interleave_samples( + dict_to_expand[key], lengths=lengths, repeat_times=expand_size + ) + elif key == "image_grid_thw": + # get the num of images for each sample + lengths = list(image_nums) + dict_to_expand[key] = _repeat_interleave_samples( + dict_to_expand[key], lengths=lengths, repeat_times=expand_size + ) + elif key == "pixel_values_videos": + samples = torch.split(video_grid_thw, list(video_nums)) + lengths = [torch.prod(sample, dim=1).sum() for sample in samples] + dict_to_expand[key] = _repeat_interleave_samples( + dict_to_expand[key], lengths=lengths, repeat_times=expand_size + ) + elif key == "video_grid_thw": + lengths = list(video_nums) + dict_to_expand[key] = _repeat_interleave_samples( + dict_to_expand[key], lengths=lengths, repeat_times=expand_size + ) + elif key == "second_per_grid_ts": + dict_to_expand[key] = _repeat_interleave_samples( + dict_to_expand[key], + lengths=list(video_nums), + repeat_times=expand_size, + ) + return dict_to_expand + + def _expand_dict_for_generation(dict_to_expand): + for key in dict_to_expand: + if ( + key != "cache_position" + and dict_to_expand[key] is not None + and isinstance(dict_to_expand[key], torch.Tensor) + and key not in visual_keys + ): + dict_to_expand[key] = dict_to_expand[key].repeat_interleave( + expand_size, dim=0 + ) + return dict_to_expand + + model_kwargs = _expand_dict_for_generation_visual(model_kwargs) + + if input_ids is not None: + input_ids = input_ids.repeat_interleave(expand_size, dim=0) + + model_kwargs = _expand_dict_for_generation(model_kwargs) + + if is_encoder_decoder: + if model_kwargs.get("encoder_outputs") is None: + raise ValueError( + "If `is_encoder_decoder` is True, make sure that `encoder_outputs` is defined." + ) + model_kwargs["encoder_outputs"] = _expand_dict_for_generation( + model_kwargs["encoder_outputs"] + ) + + return input_ids, model_kwargs diff --git a/extensions_built_in/sd_trainer/SDTrainer.py b/extensions_built_in/sd_trainer/SDTrainer.py index 63c7940..46e1f3c 100644 --- a/extensions_built_in/sd_trainer/SDTrainer.py +++ b/extensions_built_in/sd_trainer/SDTrainer.py @@ -916,7 +916,17 @@ class SDTrainer(BaseSDTrainProcess): loss = loss + norm_std_loss - return loss + additional_loss + loss = loss + additional_loss + + if self.train_config.max_loss_debug and self.train_config.max_loss is not None: + if loss.item() > self.train_config.max_loss: + print_acc(f"Loss {loss.item()} is greater than max loss {self.train_config.max_loss}. Clipping to max loss.") + print_acc(f"timesteps: {timesteps}") + + if self.train_config.max_loss is not None: + loss = torch.clamp(loss, max=self.train_config.max_loss) + + return loss def preprocess_batch(self, batch: 'DataLoaderBatchDTO'): return batch diff --git a/toolkit/advanced_prompt_embeds.py b/toolkit/advanced_prompt_embeds.py index 437ce64..b86adeb 100644 --- a/toolkit/advanced_prompt_embeds.py +++ b/toolkit/advanced_prompt_embeds.py @@ -1,4 +1,5 @@ import os +import torch from safetensors.torch import load_file, save_file @@ -30,11 +31,20 @@ class AdvancedPromptEmbeds: def __init__(self, **kwargs): self._store = {} + self._frozen_dtype_keys = [] for key, value in kwargs.items(): if not isinstance(value, list): value = [value] self._store[key] = value + @property + def frozen_dtype_keys(self): + return self._frozen_dtype_keys + + @frozen_dtype_keys.setter + def frozen_dtype_keys(self, keys): + self._frozen_dtype_keys = list(keys) if keys else [] + def __getattr__(self, name): if name.startswith("_"): raise AttributeError(name) @@ -46,10 +56,14 @@ class AdvancedPromptEmbeds: def __setattr__(self, name, value): if name.startswith("_"): super().__setattr__(name, value) - else: - if not isinstance(value, list): - value = [value] - self._store[name] = value + return + cls_attr = getattr(type(self), name, None) + if isinstance(cls_attr, property): + super().__setattr__(name, value) + return + if not isinstance(value, list): + value = [value] + self._store[name] = value def set(self, key, value): if not isinstance(value, list): @@ -77,25 +91,38 @@ class AdvancedPromptEmbeds: return key in self._store def to(self, *args, **kwargs): + frozen = set(self._frozen_dtype_keys) + if frozen: + no_dtype_args = [a for a in args if not isinstance(a, torch.dtype)] + no_dtype_kwargs = {k: v for k, v in kwargs.items() if k != "dtype"} new_pe = AdvancedPromptEmbeds() + new_pe._frozen_dtype_keys = list(self._frozen_dtype_keys) for key, value in self._store.items(): - new_pe._store[key] = [v.to(*args, **kwargs) for v in value] + if key in frozen: + new_pe._store[key] = [ + v.to(*no_dtype_args, **no_dtype_kwargs) for v in value + ] + else: + new_pe._store[key] = [v.to(*args, **kwargs) for v in value] return new_pe def detach(self): new_pe = AdvancedPromptEmbeds() + new_pe._frozen_dtype_keys = list(self._frozen_dtype_keys) for key, value in self._store.items(): new_pe._store[key] = [v.detach() for v in value] return new_pe def clone(self): new_pe = AdvancedPromptEmbeds() + new_pe._frozen_dtype_keys = list(self._frozen_dtype_keys) for key, value in self._store.items(): new_pe._store[key] = [v.clone() for v in value] return new_pe def expand_to_batch(self, batch_size): new_pe = AdvancedPromptEmbeds() + new_pe._frozen_dtype_keys = list(self._frozen_dtype_keys) for key, value in self._store.items(): if len(value) == 1: new_pe._store[key] = value * batch_size @@ -137,12 +164,18 @@ class AdvancedPromptEmbeds: cls, prompt_embeds: list["AdvancedPromptEmbeds"], padding_side: str = "right" ): embeds = {} + frozen = [] for pe in prompt_embeds: for key in pe.keys(): if key not in embeds: embeds[key] = [] embeds[key].extend(pe[key]) - return cls(**embeds) + for k in pe.frozen_dtype_keys: + if k not in frozen: + frozen.append(k) + out = cls(**embeds) + out.frozen_dtype_keys = frozen + return out @classmethod def split_prompt_embeds(cls, concatenated: "AdvancedPromptEmbeds", num_parts=None): @@ -150,6 +183,8 @@ class AdvancedPromptEmbeds: # use length of first item as num_parts num_parts = len(concatenated[concatenated.keys()[0]]) split_embeds = [cls() for _ in range(num_parts)] + for pe in split_embeds: + pe.frozen_dtype_keys = list(concatenated.frozen_dtype_keys) for key in concatenated.keys(): values = concatenated[key] if len(values) != num_parts: diff --git a/toolkit/config_modules.py b/toolkit/config_modules.py index 5ed490e..b2261c5 100644 --- a/toolkit/config_modules.py +++ b/toolkit/config_modules.py @@ -577,6 +577,11 @@ class TrainConfig: self.do_blank_stabilization = kwargs.get('do_blank_stabilization', False) self.audio_loss_multiplier = kwargs.get("audio_loss_multiplier", 1.0) + + # will throw detailed error when it goes over + self.max_loss_debug: bool = kwargs.get("max_loss_debug", False) + # will clip the loss to this amount to prevent wild outliers + self.max_loss: Optional[float] = kwargs.get("max_loss", None) ModelArch = Literal['sd1', 'sd2', 'sd3', 'sdxl', 'pixart', 'pixart_sigma', 'auraflow', 'flux', 'flex1', 'flex2', 'lumina2', 'vega', 'ssd', 'wan21'] diff --git a/ui/src/app/jobs/new/SimpleJob.tsx b/ui/src/app/jobs/new/SimpleJob.tsx index 05ff846..1f9015f 100644 --- a/ui/src/app/jobs/new/SimpleJob.tsx +++ b/ui/src/app/jobs/new/SimpleJob.tsx @@ -380,20 +380,24 @@ export default function SimpleJob({ }} options={transformerQuantizationOptions} /> - { - if (value === '') { - setJobConfig(false, 'config.process[0].model.quantize_te'); - value = defaultQtype; - } else { - setJobConfig(true, 'config.process[0].model.quantize_te'); + {!disableSections.includes('model.quantize_te') && ( + + onChange={value => { + if (value === '') { + setJobConfig(false, 'config.process[0].model.quantize_te'); + value = defaultQtype; + } else { + setJobConfig(true, 'config.process[0].model.quantize_te'); + } + setJobConfig(value, 'config.process[0].model.qtype_te'); + }} + options={quantizationOptions} + /> + )} )} {modelArch?.additionalSections?.includes('model.multistage') && ( @@ -1068,7 +1072,7 @@ export default function SimpleJob({
{[ [256, 512, 768, 1024], - [1280, 1328, 1536], + [1280, 1328, 1536, 2048], ].map(resGroup => (
{resGroup.map(res => ( diff --git a/ui/src/app/jobs/new/options.ts b/ui/src/app/jobs/new/options.ts index d35992a..7238a9c 100644 --- a/ui/src/app/jobs/new/options.ts +++ b/ui/src/app/jobs/new/options.ts @@ -6,6 +6,7 @@ type Control = 'depth' | 'line' | 'pose' | 'inpaint'; type DisableableSections = | 'model.quantize' + | 'model.quantize_te' | 'train.timestep_type' | 'network.conv' | 'trigger_word' @@ -922,6 +923,40 @@ export const modelArchs: ModelArch[] = [ disableSections: ['network.conv'], additionalSections: ['model.low_vram'], }, + { + name: 'hidream_o1', + label: 'HiDream-O1', + group: 'image', + defaults: { + 'config.process[0].model.name_or_path': ['HiDream-ai/HiDream-O1-Image', defaultNameOrPath], + 'config.process[0].model.quantize': [true, false], + 'config.process[0].model.quantize_te': [false, false], + 'config.process[0].train.timestep_type': ['linear', 'sigmoid'], + 'config.process[0].network.conv': [undefined, 16], + 'config.process[0].network.conv_alpha': [undefined, 16], + 'config.process[0].train.max_loss': [1.0, undefined], + 'config.process[0].network.network_kwargs.ignore_if_contains': [['lm_head','patch_embed', 'visual'], []], + 'config.process[0].network.transformer_only': [false, undefined], + 'config.process[0].sample.width': [2048, 1024], + 'config.process[0].sample.height': [2048, 1024], + 'config.process[0].model.model_kwargs': [ + { + noise_scale_inference: 8.0, + noise_scale: 8.0, + }, + {}, + ], + }, + disableSections: [ + 'network.conv', + 'model.quantize_te', + 'train.unload_text_encoder', + ], + additionalSections: [ + 'model.low_vram', + 'model.layer_offloading', + ], + }, ].sort((a, b) => { // Sort by label, case-insensitive return a.label.localeCompare(b.label, undefined, { sensitivity: 'base' }); diff --git a/ui/src/types.ts b/ui/src/types.ts index 087362a..d838a38 100644 --- a/ui/src/types.ts +++ b/ui/src/types.ts @@ -71,6 +71,7 @@ export interface NetworkConfig { network_kwargs: { ignore_if_contains: string[]; }; + transformer_only?: boolean; } export interface SaveConfig { @@ -149,6 +150,7 @@ export interface TrainConfig { do_differential_guidance?: boolean; differential_guidance_scale?: number; audio_loss_multiplier?: number; + max_loss?: number | null; } export interface QuantizeKwargsConfig { diff --git a/version.py b/version.py index 70e84c5..bc5433b 100644 --- a/version.py +++ b/version.py @@ -1 +1 @@ -VERSION = "0.9.8" +VERSION = "0.9.10"