diff --git a/.github/copilot-instructions.md b/.github/copilot-instructions.md index 324dd28a4..9e5476515 100644 --- a/.github/copilot-instructions.md +++ b/.github/copilot-instructions.md @@ -78,6 +78,10 @@ Use these repo-local skills for recurring SD.Next model integration work: File: `.github/skills/port-model/SKILL.md` Use when adding a new model family, porting a standalone script into a Diffusers pipeline, or wiring an upstream Diffusers model into SD.Next. +- `port-pipeline` + File: `.github/skills/port-pipeline/SKILL.md` + Use when porting a custom model pipeline implementation to a Diffusers pipeline class with behavior parity and no hard-coded device or attention assumptions. + - `debug-model` File: `.github/skills/debug-model/SKILL.md` Use when a new or existing SD.Next/Diffusers model integration fails during detection, loading, prompt encoding, sampling, or output handling. diff --git a/.github/skills/README.md b/.github/skills/README.md index a1f4a3c70..5c08ec501 100644 --- a/.github/skills/README.md +++ b/.github/skills/README.md @@ -8,6 +8,10 @@ This folder contains repo-local Copilot skills for recurring SD.Next tasks. File: `port-model/SKILL.md` Use when adding or porting a model family into SD.Next and Diffusers. +- `port-pipeline` + File: `port-pipeline/SKILL.md` + Use when porting a custom pipeline implementation into a Diffusers pipeline class while preserving behavior and avoiding hard-coded runtime assumptions. + - `debug-model` File: `debug-model/SKILL.md` Use when a new or existing SD.Next/Diffusers model integration fails during detect, load, prompt encode, sample, or output handling. diff --git a/.github/skills/port-pipeline/SKILL.md b/.github/skills/port-pipeline/SKILL.md new file mode 100644 index 000000000..2b18599c9 --- /dev/null +++ b/.github/skills/port-pipeline/SKILL.md @@ -0,0 +1,102 @@ +--- +name: port-pipeline +description: "Port custom model pipeline implementations to Diffusers. Use when migrating custom or non-Diffusers pipeline code into SD.Next repo-local pipeline files such as pipelines/model_.py or pipelines//pipeline.py while preserving behavior, avoiding new dependencies, and keeping device/attention handling configurable." +argument-hint: "Provide source pipeline path, target SD.Next destination path, and target pipeline class name" +--- + +# Port Custom Pipeline To Diffusers + +Port an existing custom model pipeline implementation into a Diffusers-compatible pipeline class with behavior parity and SD.Next-friendly conventions. +This skill targets SD.Next repo-local pipeline ports only. + +## When To Use + +- A user has a custom pipeline implementation and wants it ported to Diffusers +- Existing model code is runnable but not structured as a Diffusers pipeline +- The destination is SD.Next pipeline code under `pipelines/model_*.py` or `pipelines//` +- The task requires preserving generation behavior without introducing new dependencies +- The task requires removing hard-coded runtime assumptions (device or attention backend) + +## Mandatory Clarification Gate + +Before implementation, confirm these required inputs with the user: + +1. Path to the source custom pipeline implementation +2. Destination path in this SD.Next repository (typically under `pipelines/`) +3. Target pipeline class name + +If any of the above are missing or ambiguous, stop and ask concise clarification questions before writing code. + +## Constraints + +- Do not add new dependencies +- Do not hard-code device type (`cpu`, `cuda`, `mps`, etc.) +- Do not hard-code attention type or backend assumptions +- Preserve externally visible behavior of the source pipeline unless the user asks for intentional changes + +## Workflow + +1. Collect Inputs +- Ask for source path, destination path, and target pipeline name. +- Confirm destination is an SD.Next repo-local pipeline location, not an upstream Diffusers repository path. +- Confirm runtime assumptions and expected task type (text-to-image, image-to-image, inpaint, etc.). + +2. Analyze Source Pipeline +- Inspect model loading, prompt processing, denoising or sampling loop, scheduler interactions, and output post-processing. +- Identify all components that must be ported: models, tokenizers or processors, schedulers, adapters, preprocessors, postprocessors, callbacks, and output dataclasses. +- Note any hidden global state, side effects, or implicit defaults that must become explicit parameters. + +3. Map To Diffusers Interfaces +- Choose the most appropriate Diffusers base class and output type. +- Define `__init__`, module registration, `from_pretrained` and `__call__` signatures aligned with existing Diffusers patterns. +- Keep parameter names and behavior as close as possible to upstream conventions. +- Identify any custom classes needed beyond the pipeline itself: transformer blocks, attention processors, custom schedulers, or output types. Plan a separate module file for each. + +4. Implement Supporting Classes +- If the pipeline requires custom model classes (e.g., a custom transformer block, attention module, or other model component), implement each in a **separate module** located in the **same directory** as the main pipeline file (e.g., `pipelines//transformer.py`, `pipelines//scheduler.py`). +- If the pipeline requires a custom scheduler class, implement it in its own module (e.g., `pipelines//scheduler_.py`) following Diffusers scheduler conventions (`step`, `add_noise`, `scale_model_input`, etc.). +- Each supporting class module must be self-contained: no circular imports, no hidden global state, and no hard-coded device or attention assumptions. +- Import supporting classes into the main pipeline module from their respective sibling modules. + +5. Implement Pipeline Class +- Create the destination pipeline classes at the user-provided path. +- Port logic in small, testable sections: initialization, input validation, prompt encoding, latent preparation, denoising loop, decoding, and output packaging. +- Replace hard-coded device and attention logic with runtime-configurable behavior. +- Keep imports limited to existing project and Diffusers dependencies. + +6. Lint And Fix +- Activate the project venv: `source venv/bin/activate` +- Run `ruff` on all newly written files: `pnpm ruff` (or `ruff check --fix` for targeted runs). +- Run `pylint` on all newly written files: `pnpm pylint` (or `pylint ` for targeted runs). +- Fix every reported error or warning that is not explicitly marked with a `TODO` suppression comment in the source. +- Re-run both linters after fixes to confirm a clean result before proceeding. + +7. Validate Behavior Parity +- Compare source and ported implementations for input-output shape handling, dtype flow, scheduler step ordering, and guidance behavior. +- Run focused checks or smoke tests if available in the workspace. +- Call out any known differences that were required for Diffusers compatibility. + +8. Report Results +- Summarize what was ported and where. +- List any unresolved assumptions, risks, or TODOs. +- Provide minimal follow-up steps for integration and testing. + +## Review Checklist + +- Required inputs were collected before edits +- No new dependency was introduced +- No hard-coded device or attention backend remains +- Core components from source pipeline were fully mapped +- Pipeline class is in requested destination with requested name +- Each custom supporting class (transformer, scheduler, etc.) is in its own sibling module +- Supporting modules have no circular imports or hidden global state +- `ruff` and `pylint` both pass cleanly on all newly written files (venv activated) +- Main inference path behavior matches the source implementation + +## Output Expectations + +Final response should include: +- Source path, destination path, and final pipeline class name +- Brief parity summary of key components ported +- Validation performed and any gaps +- Explicit note of any assumptions requiring user confirmation diff --git a/CHANGELOG.md b/CHANGELOG.md index 7a8f8c77e..94fcd7505 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -32,6 +32,12 @@ For full details, see [ChangeLog](https://github.com/vladmandic/automatic/blob/m still a popular method for upscaling, but has not been updated nor maintained for a while so now its modernized and fully integrated as a built-in script! - **Models** +HiDream-O1-Image is a natively unified image generative foundation model built on a Pixel-level Unified Transformer (UiT) without external VAEs or disjoint text encoders, which natively encodes raw pixels, text, and task-specific conditions in a single shared token space — supporting text-to-image, image editing, and subject-driven personalization at up to 2,048 × 2,048. + - [HiDream-O1-Image](https://huggingface.co/HiDream-ai/HiDream-O1-Image) pixel-level unified transformer model support + HiDream-O1 is based on a single custom *Qwen3-VL* 8.8B 35GB component + includes both **HiDream-O1-Image** *(base)* and **HiDream-O1-Image-Dev** *(distilled*)* variants + includes *T2I* and *I2I edit* capabilities and resolutions up to 2048px + *note*: use steps:50 for base and steps:28 for dev variants - [JoyAI Image Edit](https://huggingface.co/jdopensource/JoyAI-Image-Edit-Diffusers) image-editing model support includes multimodal conditioning using *Qwen3-VL* with a dedicated *JoyImageEdit* diffusion transformer *note* this is a large model at 50GB so use of agressive quantization is recommended diff --git a/TODO.md b/TODO.md index e607bb7cb..0156448cb 100644 --- a/TODO.md +++ b/TODO.md @@ -4,7 +4,6 @@ ### Assigned -- Check Outpaint, @vladmandic - Chat-based interface, @vladmandic - Control tab verify overrides handling, @vladmandic - Reimplement `llama` remover for Kanvas, @vladmandic diff --git a/data/reference-distilled.json b/data/reference-distilled.json index ad7b07ddb..4588ae5c1 100644 --- a/data/reference-distilled.json +++ b/data/reference-distilled.json @@ -36,6 +36,16 @@ "skip": true, "extras": "sampler: Default, cfg_scale: 4.5" }, + "HiDream-O1 Image Dev": { + "path": "HiDream-ai/HiDream-O1-Image-Dev", + "preview": "HiDream-ai--HiDream-O1-Image-Dev.jpg", + "desc": "HiDream-O1-Image-Dev is the distilled 8B HiDream-O1 variant tuned for 28-step fast generation using flash flow scheduling.", + "skip": true, + "extras": "sampler: Flash, steps: 28, cfg_scale: 0.0", + "size": 35.2, + "tags": "distilled", + "date": "2026 May" + }, "Qwen-Image-Lightning": { "path": "vladmandic/Qwen-Lightning", "preview": "vladmandic--Qwen-Lightning.jpg", diff --git a/data/reference.json b/data/reference.json index 61ebecabf..0a00af7cf 100644 --- a/data/reference.json +++ b/data/reference.json @@ -669,6 +669,15 @@ "size": 58.4, "date": "2025 April" }, + "HiDream-O1 Image": { + "path": "HiDream-ai/HiDream-O1-Image", + "desc": "HiDream-O1-Image is an 8B pixel-level unified transformer model for text-to-image generation, instruction editing, and multi-reference personalization up to 2048x2048.", + "preview": "HiDream-ai--HiDream-O1-Image.jpg", + "skip": true, + "extras": "sampler: Default", + "size": 35.2, + "date": "2026 May" + }, "HiDream-E1 Full": { "path": "HiDream-ai/HiDream-E1-Full", "desc": "HiDream-E1 is an image editing model built on HiDream-I1.", diff --git a/models/Reference/HiDream-ai--HiDream-O1-Image-Dev.jpg b/models/Reference/HiDream-ai--HiDream-O1-Image-Dev.jpg new file mode 100644 index 000000000..e69de29bb diff --git a/models/Reference/HiDream-ai--HiDream-O1-Image.jpg b/models/Reference/HiDream-ai--HiDream-O1-Image.jpg new file mode 100644 index 000000000..e69de29bb diff --git a/modules/control/run.py b/modules/control/run.py index d3238675f..1327d9507 100644 --- a/modules/control/run.py +++ b/modules/control/run.py @@ -281,7 +281,7 @@ def control_process(p: StableDiffusionProcessingControl, input_image: Image.Image = None, # only used for tiling, otherwise processor.preprocess_image set p params ): debug_log(f'Control exec pipeline: task={sd_models.get_diffusers_task(pipe)} class={pipe.__class__}') - if sd_models.get_diffusers_task(pipe) != sd_models.DiffusersTaskType.TEXT_2_IMAGE: # force vae back to gpu if not in txt2img mode + if sd_models.get_diffusers_task(pipe) != sd_models.DiffusersTaskType.TEXT_2_IMAGE and hasattr(pipe, 'vae'): # force vae back to gpu if not in txt2img mode sd_models.move_model(pipe.vae, devices.device) # what are we doing? diff --git a/modules/modeldata.py b/modules/modeldata.py index 1c2acf2d7..4662f8165 100644 --- a/modules/modeldata.py +++ b/modules/modeldata.py @@ -72,6 +72,8 @@ def get_model_type(pipe): model_type = 'sana' elif 'VIBE' in name: model_type = 'sana' + elif "HiDreamO1" in name: + model_type = 'o1' elif "HiDream" in name: model_type = 'h1' elif name.startswith("Anima") and "AnimateDiff" not in name: diff --git a/modules/processing_callbacks.py b/modules/processing_callbacks.py index 97bfc0fc4..c9689c6e5 100644 --- a/modules/processing_callbacks.py +++ b/modules/processing_callbacks.py @@ -59,6 +59,23 @@ def diffusers_callback(pipe, step: int = 0, timestep: int = 0, kwargs: dict | No torch.xpu.synchronize(devices.device) elif devices.backend in {"cuda", "zluda", "rocm"}: torch.cuda.synchronize(devices.device) + + if shared.state.paused: + log.debug('Sampling paused') + while shared.state.paused: + if shared.state.interrupted or shared.state.skipped: + raise AssertionError('Interrupted...') + time.sleep(0.1) + + image = kwargs.get('image', None) + if image is not None: + shared.state.current_image = image + shared.state.current_latent = None + shared.state.step() # increase step + shared.state.preview_job = -1 # indicate that preview image has changed + debug_callback(f'Callback: step={step} timestep={timestep} image={image if image is not None else None} kwargs={list(kwargs)}') + return kwargs + latents = kwargs.get('latents', None) if debug: debug_callback(f'Callback: step={step} timestep={timestep} latents={latents.shape if latents is not None else None} kwargs={list(kwargs)}') @@ -67,12 +84,6 @@ def diffusers_callback(pipe, step: int = 0, timestep: int = 0, kwargs: dict | No shared.state.step() if shared.state.interrupted or shared.state.skipped: raise AssertionError('Interrupted...') - if shared.state.paused: - log.debug('Sampling paused') - while shared.state.paused: - if shared.state.interrupted or shared.state.skipped: - raise AssertionError('Interrupted...') - time.sleep(0.1) if latents is None: return kwargs elif shared.opts.nan_skip: diff --git a/modules/processing_correction.py b/modules/processing_correction.py index 1c3751bb9..9477960dd 100644 --- a/modules/processing_correction.py +++ b/modules/processing_correction.py @@ -212,6 +212,8 @@ def _count_steps_below(pipe, threshold): def correction_callback(p, timestep, kwargs, pipe=None, initial: bool = False, step: int = 0): + if pipe and pipe.__class__.__name__ in ['HiDreamO1Pipeline', 'HiDreamO1ImagePipeline']: + return kwargs if initial: if not any([p.hdr_clamp, p.hdr_mode, p.hdr_maximize, p.hdr_sharpen, p.hdr_color, p.hdr_brightness, p.hdr_tint_ratio]): p.correction_skip = True diff --git a/modules/schedulers/scheduler_flashflow.py b/modules/schedulers/scheduler_flashflow.py index e9a82c952..f7df144d2 100644 --- a/modules/schedulers/scheduler_flashflow.py +++ b/modules/schedulers/scheduler_flashflow.py @@ -184,8 +184,8 @@ class FlashFlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin): def set_timesteps( self, - num_inference_steps: int = None, - device: Union[str, torch.device] = None, + num_inference_steps: Optional[int] = None, + device: Optional[Union[str, torch.device]] = None, sigmas: Optional[List[float]] = None, mu: Optional[float] = None, ): @@ -288,6 +288,7 @@ class FlashFlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin): s_tmin: float = 0.0, s_tmax: float = float("inf"), s_noise: float = 1.0, + noise_clip_std: float = 0.0, generator: Optional[torch.Generator] = None, return_dict: bool = True, ) -> Union[FlashFlowMatchEulerDiscreteSchedulerOutput, Tuple]: @@ -352,7 +353,9 @@ class FlashFlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin): device=model_output.device, dtype=denoised.dtype, ) - sample = sigma_next * noise + (1.0 - sigma_next) * denoised + if noise_clip_std > 0.0: + noise = noise.clamp(-noise_clip_std, noise_clip_std) + sample = sigma_next * s_noise * noise + (1.0 - sigma_next) * denoised self._step_index += 1 sample = sample.to(model_output.dtype) diff --git a/modules/schedulers/scheduler_unipc_flowmatch.py b/modules/schedulers/scheduler_unipc_flowmatch.py index bea747373..1adbd0799 100644 --- a/modules/schedulers/scheduler_unipc_flowmatch.py +++ b/modules/schedulers/scheduler_unipc_flowmatch.py @@ -654,8 +654,14 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin): model_output: torch.Tensor, timestep: Union[int, torch.Tensor], sample: torch.Tensor, + s_churn: float = 0.0, + s_tmin: float = 0.0, + s_tmax: float = float("inf"), + s_noise: float = 1.0, + noise_clip_std: float = 0.0, return_dict: bool = True, - generator=None) -> Union[SchedulerOutput, Tuple]: + generator=None, + **kwargs) -> Union[SchedulerOutput, Tuple]: """ Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with the multistep UniPC. diff --git a/modules/sd_detect.py b/modules/sd_detect.py index 12acc4f85..a3978f002 100644 --- a/modules/sd_detect.py +++ b/modules/sd_detect.py @@ -87,6 +87,10 @@ def guess_by_name(fn, current_guess): new_guess = 'OmniGen' elif 'sd3' in fn.lower(): new_guess = 'Stable Diffusion 3' + elif 'hidream-o1' in fn.lower(): + new_guess = 'HiDreamO1' + elif 'hidream' in fn.lower(): + new_guess = 'HiDream' elif 'hidream' in fn.lower(): new_guess = 'HiDream' elif 'zeta-chroma' in fn.lower() or 'zetachroma' in fn.lower(): diff --git a/modules/sd_models.py b/modules/sd_models.py index fb1f6e7a3..288d10ab3 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -433,6 +433,10 @@ def load_diffuser_force(detected_model_type, checkpoint_info, diffusers_load_con from pipelines.model_omnigen import load_omnigen sd_model = load_omnigen(checkpoint_info, diffusers_load_config) allow_post_quant = False + elif model_type in ['HiDreamO1']: + from pipelines.model_hidream import load_hidream_o1 + sd_model = load_hidream_o1(checkpoint_info, diffusers_load_config) + allow_post_quant = False elif model_type in ['HiDream']: from pipelines.model_hidream import load_hidream sd_model = load_hidream(checkpoint_info, diffusers_load_config) @@ -1285,6 +1289,8 @@ def set_diffuser_pipe(pipe, new_pipe_type): fn = f'{sys._getframe(2).f_code.co_name}:{sys._getframe(1).f_code.co_name}' # pylint: disable=protected-access log.trace(f"Pipeline class change requested: target={new_pipe_type} fn={fn}") # pylint: disable=protected-access log.warning(f'Pipeline class change failed: type={new_pipe_type} pipeline={cls} {e}') + if debug_load: + errors.display(e, 'Pipeline switch') has_errors = True if not hasattr(pipe, 'config') or has_errors: try: # maybe a wrapper pipeline so just change the class @@ -1302,6 +1308,8 @@ def set_diffuser_pipe(pipe, new_pipe_type): return pipe except Exception as e: # pylint: disable=unused-variable log.warning(f'Pipeline class set failed: type={new_pipe_type} pipeline={cls} {e}') + if debug_load: + errors.display(e, 'Pipeline switch') has_errors = True return pipe diff --git a/modules/shared_state.py b/modules/shared_state.py index 4bfe6fae5..d49310c26 100644 --- a/modules/shared_state.py +++ b/modules/shared_state.py @@ -270,7 +270,7 @@ class State: def do_set_current_image(self): from modules import shared, images, sd_samplers_common - if (self.current_latent is None) or self.disable_preview or (self.preview_job == self.job_no): + if self.disable_preview or (self.preview_job == self.job_no): return False self.preview_job = self.job_no @@ -280,27 +280,33 @@ class State: self.preview_job = -1 return True - try: - sample = self.current_latent - self.current_image_sampling_step = self.sampling_step + if self.current_latent is not None: try: - if self.current_noise_pred is not None and self.current_sigma is not None and self.current_sigma_next is not None: - original_sample = sample - (self.current_noise_pred * (self.current_sigma_next-self.current_sigma)) - if self.prediction_type in {"epsilon", "flow_prediction"}: - sample = original_sample - (self.current_noise_pred * self.current_sigma) - elif self.prediction_type == "v_prediction": - sample = self.current_noise_pred * (-self.current_sigma / (self.current_sigma**2 + 1) ** 0.5) + (original_sample / (self.current_sigma**2 + 1)) # pylint: disable=invalid-unary-operand-type - except Exception: - pass # ignore sigma errors - image = sd_samplers_common.samples_to_image_grid(sample) - self.assign_current_image(image) + sample = self.current_latent + self.current_image_sampling_step = self.sampling_step + try: + if self.current_noise_pred is not None and self.current_sigma is not None and self.current_sigma_next is not None: + original_sample = sample - (self.current_noise_pred * (self.current_sigma_next-self.current_sigma)) + if self.prediction_type in {"epsilon", "flow_prediction"}: + sample = original_sample - (self.current_noise_pred * self.current_sigma) + elif self.prediction_type == "v_prediction": + sample = self.current_noise_pred * (-self.current_sigma / (self.current_sigma**2 + 1) ** 0.5) + (original_sample / (self.current_sigma**2 + 1)) # pylint: disable=invalid-unary-operand-type + except Exception: + pass # ignore sigma errors + image = sd_samplers_common.samples_to_image_grid(sample) + self.assign_current_image(image) + self.preview_job = -1 + return True + except Exception as e: + self.preview_job = -1 + log.error(f'State image: last={self.id_live_preview} step={self.sampling_step} {e}') + display(e, 'State image') + return False + elif self.current_image is not None: + self.assign_current_image(self.current_image) self.preview_job = -1 return True - except Exception as e: - self.preview_job = -1 - log.error(f'State image: last={self.id_live_preview} step={self.sampling_step} {e}') - display(e, 'State image') - return False + return False def assign_current_image(self, image): self.current_image = image diff --git a/pipelines/hidream/pipeline_hidream_image_editing.py b/pipelines/hidream/hidream_e1.py similarity index 100% rename from pipelines/hidream/pipeline_hidream_image_editing.py rename to pipelines/hidream/hidream_e1.py diff --git a/pipelines/hidream/hidream_o1.py b/pipelines/hidream/hidream_o1.py new file mode 100644 index 000000000..5000f48de --- /dev/null +++ b/pipelines/hidream/hidream_o1.py @@ -0,0 +1,458 @@ +from dataclasses import dataclass +from typing import Any, Callable, Dict, List, Optional, Union + +import numpy as np +import PIL.Image +import torch +from tqdm.rich import tqdm + +from diffusers.pipelines.pipeline_utils import DiffusionPipeline +from diffusers.utils import BaseOutput +from diffusers.utils.torch_utils import randn_tensor + +from pipelines.hidream.scheduler_flashfloweuler import FlashFlowMatchEulerDiscreteScheduler +from pipelines.hidream.scheduler_flowunipc import FlowUniPCMultistepScheduler +from pipelines.hidream.hidream_o1_utils import ( + CONDITION_IMAGE_SIZE, + DEFAULT_TIMESTEPS, + NOISE_SCALE, + PATCH_SIZE, + TIMESTEP_TOKEN_NUM, + T_EPS, + _calculate_dimensions, + _ensure_special_tokens, + _image_to_patch_tensor, + _patches_to_np, + _pil_to_normalized_tensor, + _resize_pilimage, + build_t2i_text_sample, + get_rope_index_fix_point, +) + +use_flash_attn = False +try: + import flash_attn # pylint: disable=unused-import + use_flash_attn = True +except ImportError: + pass + + +@dataclass +class HiDreamO1PipelineOutput(BaseOutput): + images: Union[List[PIL.Image.Image], np.ndarray] + + +class HiDreamO1Pipeline(DiffusionPipeline): + model_cpu_offload_seq = "transformer" + _callback_tensor_inputs = ["latents"] + vae_scale_factor = 1 + + def __init__( + self, + transformer, + processor, + tokenizer, + scheduler, + ): + super().__init__() + self.register_modules(transformer=transformer, processor=processor, scheduler=scheduler, tokenizer=tokenizer) + _ensure_special_tokens(self.tokenizer) + + def _build_scheduler( + self, + num_inference_steps: int, + shift: float, + device: torch.device, + ): + if num_inference_steps <= 28: + self.scheduler = FlashFlowMatchEulerDiscreteScheduler(num_train_timesteps=1000, shift=shift, use_dynamic_shifting=False) + timesteps_list = DEFAULT_TIMESTEPS if num_inference_steps == 28 else None + else: + self.scheduler = FlowUniPCMultistepScheduler(use_dynamic_shifting=False, shift=shift) + timesteps_list = None + self.scheduler.set_timesteps(num_inference_steps, device=device) + if timesteps_list is not None: + self.scheduler.timesteps = torch.tensor(timesteps_list, device=device, dtype=torch.long) + sigmas = [t.item() / 1000.0 for t in self.scheduler.timesteps] + sigmas.append(0.0) + self.scheduler.sigmas = torch.tensor(sigmas, device=device) + + def _prepare_reference_paths( + self, + image: Optional[Union[PIL.Image.Image, List[PIL.Image.Image]]], + ref_images: Optional[Union[PIL.Image.Image, List[PIL.Image.Image]]], + ) -> List[PIL.Image.Image]: + refs: List[PIL.Image.Image] = [] + if image is not None: + if isinstance(image, list): + refs.extend(image) + else: + refs.append(image) + if ref_images is not None: + if isinstance(ref_images, list): + refs.extend(ref_images) + else: + refs.append(ref_images) + return refs + + @torch.no_grad() + def __call__( + self, + prompt: str, + negative_prompt: Optional[str] = None, + image: Optional[Union[PIL.Image.Image, List[PIL.Image.Image]]] = None, + ref_images: Optional[Union[PIL.Image.Image, List[PIL.Image.Image]]] = None, + height: int = 1440, + width: int = 2560, + num_inference_steps: int = 50, + guidance_scale: float = 5.0, + shift: float = 3.0, + timesteps_list: Optional[List[int]] = None, + scheduler: Optional[Union[FlashFlowMatchEulerDiscreteScheduler, FlowUniPCMultistepScheduler]] = None, + generator: Optional[torch.Generator] = None, + seed: Optional[int] = None, + noise_scale_start: float = NOISE_SCALE, + noise_scale_end: float = NOISE_SCALE, + noise_clip_std: float = 0.0, + keep_original_aspect: bool = True, + callback_on_step_end: Optional[Callable[[DiffusionPipeline, int, int, Dict[str, torch.Tensor]], Dict[str, torch.Tensor]]] = None, + callback_on_step_end_tensor_inputs: Optional[List[str]] = None, + callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, + callback_steps: int = 1, + output_type: str = "pil", + return_dict: bool = True, + **kwargs, + ) -> Union[HiDreamO1PipelineOutput, tuple]: + model = self.transformer + processor = self.processor + tokenizer = self.tokenizer + model_config = model.config + + if isinstance(prompt, list): + prompt = prompt[0] if len(prompt) > 0 else "" + if isinstance(negative_prompt, list): + negative_prompt = negative_prompt[0] if len(negative_prompt) > 0 else "" + if num_inference_steps <= 28: + guidance_scale = 1.0 + + device = self._execution_device + try: + dtype = next(model.parameters()).dtype + except (StopIteration, AttributeError, TypeError): + dtype = torch.bfloat16 + + if callback_on_step_end_tensor_inputs is not None: + invalid_inputs = [name for name in callback_on_step_end_tensor_inputs if name not in self._callback_tensor_inputs] + if invalid_inputs: + raise ValueError( + f"callback_on_step_end_tensor_inputs has to be in {self._callback_tensor_inputs}, but found {invalid_inputs}" + ) + + refs = [img.convert("RGB") for img in self._prepare_reference_paths(image, ref_images)] + preresized_ref_pil = None + + if keep_original_aspect and len(refs) >= 1: + preresized_ref_pil = _resize_pilimage(refs[0], 2048, PATCH_SIZE) + width, height = preresized_ref_pil.size + else: + width = max(PATCH_SIZE, int(round(width / PATCH_SIZE)) * PATCH_SIZE) + height = max(PATCH_SIZE, int(round(height / PATCH_SIZE)) * PATCH_SIZE) + + h_patches = height // PATCH_SIZE + w_patches = width // PATCH_SIZE + + if len(refs) == 0: + cond_sample = build_t2i_text_sample(prompt, height, width, tokenizer, processor, model_config) + uncond_sample = None + if guidance_scale > 1.0: + uncond_sample = build_t2i_text_sample(negative_prompt or " ", height, width, tokenizer, processor, model_config) + + def to_device(sample): + return {k: (v.to(device) if torch.is_tensor(v) else v) for k, v in sample.items()} + + cond_sample = to_device(cond_sample) + if uncond_sample is not None: + uncond_sample = to_device(uncond_sample) + + ref_patches = None + tgt_image_len = (height // PATCH_SIZE) * (width // PATCH_SIZE) + samples = [cond_sample] + if uncond_sample is not None: + samples.append(uncond_sample) + else: + 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 + spatial_merge_size = model_config.vision_config.spatial_merge_size + + ref_pils = [preresized_ref_pil] if preresized_ref_pil is not None else refs + k_refs = len(ref_pils) + + max_size = max(height, width) + if k_refs == 2: + max_size = max_size * 48 // 64 + elif k_refs <= 4: + max_size = max_size // 2 + elif k_refs <= 8: + max_size = max_size * 24 // 64 + elif k_refs > 8: + max_size = max_size // 4 + + ref_pils_resized, ref_patches_list = [], [] + for pil in ref_pils: + pil_r = pil if (preresized_ref_pil is not None and pil is preresized_ref_pil) else _resize_pilimage(pil, max_size, PATCH_SIZE) + ref_pils_resized.append(pil_r) + x = _pil_to_normalized_tensor(pil_r).unsqueeze(0) + x = _image_to_patch_tensor(x, patch_size=PATCH_SIZE).squeeze(0) + ref_patches_list.append(x) + + ref_image_lens = [img.shape[0] for img in ref_patches_list] + total_ref_len = sum(ref_image_lens) + ref_patches = torch.cat(ref_patches_list, dim=0).unsqueeze(0).to(device, dtype) + + tgt_image_len = (height // PATCH_SIZE) * (width // PATCH_SIZE) + + cond_img_size = CONDITION_IMAGE_SIZE + if k_refs > 4 and k_refs <= 8: + cond_img_size = CONDITION_IMAGE_SIZE * 48 // 64 + elif k_refs > 8: + cond_img_size = CONDITION_IMAGE_SIZE // 2 + + ref_pils_vlm = [] + for pil_r in ref_pils_resized: + cond_w, cond_h = _calculate_dimensions(cond_img_size, pil_r.width / pil_r.height) + ref_pils_vlm.append(pil_r.resize((cond_w, cond_h), resample=PIL.Image.Resampling.LANCZOS)) + + image_grid_thw_tgt = torch.tensor([1, height // PATCH_SIZE, width // PATCH_SIZE], dtype=torch.int64).unsqueeze(0) + image_grid_thw_ref = torch.zeros((k_refs, 3), dtype=torch.int64) + for i, pil_r in enumerate(ref_pils_resized): + rw, rh = pil_r.size + image_grid_thw_ref[i] = torch.tensor([1, rh // PATCH_SIZE, rw // PATCH_SIZE], dtype=torch.int64) + + samples = [] + captions = [prompt] + if guidance_scale > 1.0: + captions.append(negative_prompt or " ") + + for caption in captions: + boi_token = getattr(tokenizer, "boi_token", "<|boi_token|>") + tms_token = getattr(tokenizer, "tms_token", "<|tms_token|>") + + content = [{"type": "image"} for _ in range(k_refs)] + content.append({"type": "text", "text": caption}) + messages = [{"role": "user", "content": content}] + + template_caption = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + proc = processor(text=[template_caption], images=ref_pils_vlm, padding="longest", return_tensors="pt") + + input_ids_2 = tokenizer.encode(boi_token + tms_token * TIMESTEP_TOKEN_NUM, return_tensors="pt", add_special_tokens=False) + input_ids = torch.cat([proc.input_ids, input_ids_2], dim=-1) + + igthw_cond = proc.image_grid_thw.clone() + for i in range(k_refs): + igthw_cond[i, 1] //= spatial_merge_size + igthw_cond[i, 2] //= spatial_merge_size + igthw_all = torch.cat([igthw_cond, image_grid_thw_tgt, image_grid_thw_ref], dim=0) + + vision_tokens_list = [] + vt_tgt = torch.full((1, tgt_image_len), image_token_id, dtype=input_ids.dtype) + vt_tgt[0, 0] = vision_start_token_id + vision_tokens_list.append(vt_tgt) + for ref_len in ref_image_lens: + vt_ref = torch.full((1, ref_len), image_token_id, dtype=input_ids.dtype) + vt_ref[0, 0] = vision_start_token_id + vision_tokens_list.append(vt_ref) + vision_tokens = torch.cat(vision_tokens_list, dim=1) + input_ids_pad = torch.cat([input_ids, vision_tokens], dim=-1) + + position_ids, _ = get_rope_index_fix_point( + 1, + image_token_id, + video_token_id, + vision_start_token_id, + input_ids=input_ids_pad, + image_grid_thw=igthw_all, + video_grid_thw=None, + attention_mask=None, + skip_vision_start_token=[0] * k_refs + [1] + [1] * k_refs, + ) + + txt_seq_len = input_ids.shape[-1] + all_seq_len = position_ids.shape[-1] + + token_types_raw = torch.zeros((1, all_seq_len), dtype=input_ids.dtype) + bgn = txt_seq_len - TIMESTEP_TOKEN_NUM + end = bgn + tgt_image_len + TIMESTEP_TOKEN_NUM + token_types_raw[0, bgn:end] = 1 + token_types_raw[0, end : end + total_ref_len] = 2 + token_types_raw[0, txt_seq_len - TIMESTEP_TOKEN_NUM : txt_seq_len] = 3 + + vinput_mask = torch.logical_or(token_types_raw == 1, token_types_raw == 2) + token_types_bin = (token_types_raw > 0).to(token_types_raw.dtype) + + samples.append( + { + "input_ids": input_ids_pad.to(device), + "position_ids": position_ids.to(device), + "token_types": token_types_bin.to(device), + "vinput_mask": vinput_mask.to(device), + "pixel_values": proc.pixel_values.to(device, dtype), + "image_grid_thw": proc.image_grid_thw.to(device), + } + ) + + if generator is None: + generator = torch.Generator() + if seed is not None: + generator.manual_seed(seed + 1) + + noise = noise_scale_start * randn_tensor((1, 3, height, width), generator=generator, device=device, dtype=dtype) + z = _image_to_patch_tensor(noise, patch_size=PATCH_SIZE) + + + if scheduler is None: + self._build_scheduler( + num_inference_steps=num_inference_steps, + shift=shift, + device=device, + ) + + num_steps = len(self.scheduler.timesteps) + if num_steps > 1: + noise_scale_schedule = [noise_scale_start + (noise_scale_end - noise_scale_start) * i / (num_steps - 1) for i in range(num_steps)] + else: + noise_scale_schedule = [noise_scale_start] + + def forward_once(sample: Dict[str, torch.Tensor], z_in: torch.Tensor, t_pixeldit: torch.Tensor): + kwargs: Dict[str, Any] = { + "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"], + } + if use_flash_attn is not None: + kwargs["use_flash_attn"] = use_flash_attn + if "pixel_values" in sample: + kwargs["pixel_values"] = sample["pixel_values"] + if "image_grid_thw" in sample: + kwargs["image_grid_thw"] = sample["image_grid_thw"] + + outputs = model(**kwargs) + x_pred = outputs.x_pred + if ref_patches is None: + return x_pred[0, sample["vinput_mask"][0]].unsqueeze(0) + return x_pred[0, sample["vinput_mask"][0]][:tgt_image_len].unsqueeze(0) + + preview_x0 = None + for step_idx, step_t in enumerate(tqdm(self.scheduler.timesteps, desc="Processing", unit="it")): + t_pixeldit = 1.0 - step_t.float() / 1000.0 + sigma = (step_t.float() / 1000.0).to(dtype=torch.float32).clamp_min(T_EPS) + + if ref_patches is None: + x_pred_cond = forward_once(samples[0], z.clone(), t_pixeldit) + v_cond = (x_pred_cond.to(dtype=torch.float32) - z.to(dtype=torch.float32)) / sigma + if len(samples) > 1: + x_pred_uncond = forward_once(samples[1], z.clone(), t_pixeldit) + v_uncond = (x_pred_uncond.to(dtype=torch.float32) - z.to(dtype=torch.float32)) / sigma + v_guided = v_uncond + guidance_scale * (v_cond - v_uncond) + else: + v_guided = v_cond + preview_x0 = x_pred_cond + else: + vinputs = torch.cat([z, ref_patches], dim=1) + x_vis_list = [forward_once(sample, vinputs, t_pixeldit) for sample in samples] + x_vis_stacked = torch.cat(x_vis_list, dim=0) + + z_rep = z.expand(len(samples), -1, -1) + v_pred = (x_vis_stacked.to(dtype=torch.float32) - z_rep.to(dtype=torch.float32)) / sigma + v_cond = v_pred[0:1] + if len(samples) > 1: + v_uncond = v_pred[1:] + v_guided = v_uncond + guidance_scale * (v_cond - v_uncond) + else: + v_guided = v_cond + preview_x0 = x_vis_list[0] + + model_output = -v_guided + if num_inference_steps <= 28: + z = self.scheduler.step( + model_output.float(), + step_t.to(dtype=torch.float32), + z.float(), + s_noise=noise_scale_schedule[step_idx], + noise_clip_std=noise_clip_std, + return_dict=False, + )[0].to(dtype) + else: + z = self.scheduler.step(model_output.float(), step_t.to(dtype=torch.float32), z.float(), return_dict=False)[0].to(dtype) + + if callback_on_step_end is not None: + pil_image = PIL.Image.fromarray(_patches_to_np(preview_x0, h_patches, w_patches, invert=False)) + callback_on_step_end(self, step_idx, int(step_t.item()), {"image": pil_image}) + + if output_type == "pil": + output_images = [PIL.Image.fromarray(_patches_to_np(preview_x0, h_patches, w_patches, invert=False, rescale=False))] + elif output_type == "np": + output_images = [_patches_to_np(preview_x0, h_patches, w_patches, invert=True, rescale=True)] + else: + raise ValueError(f"Unsupported output_type={output_type!r}; supported values are 'pil' and 'np'") + + if not return_dict: + return (output_images,) + return HiDreamO1PipelineOutput(images=output_images) + + +class HiDreamO1ImagePipeline(HiDreamO1Pipeline): + def __call__( + self, + prompt: str, + negative_prompt: Optional[str] = None, + image: Optional[Union[PIL.Image.Image, List[PIL.Image.Image]]] = None, + height: int = 1440, + width: int = 2560, + num_inference_steps: int = 50, + guidance_scale: float = 5.0, + shift: float = 3.0, + timesteps_list: Optional[List[int]] = None, + scheduler: Optional[Union[FlashFlowMatchEulerDiscreteScheduler, FlowUniPCMultistepScheduler]] = None, + generator: Optional[torch.Generator] = None, + seed: Optional[int] = None, + noise_scale_start: float = NOISE_SCALE, + noise_scale_end: float = NOISE_SCALE, + noise_clip_std: float = 0.0, + keep_original_aspect: bool = True, + use_flash_attn: Optional[bool] = None, + callback: Optional[Callable[[int, int, Callable[[], PIL.Image.Image]], None]] = None, + output_type: str = "pil", + return_dict: bool = True, + **kwargs, + ) -> Union[HiDreamO1PipelineOutput, tuple]: + # image is list, first entry should go to image and remaining to ref_images + ref_images = [] + if isinstance(image, list): + ref_images = image[1:] + image = image[0] if len(image) > 0 else None + return super().__call__( + prompt=prompt, + negative_prompt=negative_prompt, + image=image, + ref_images=ref_images, + height=height, + width=width, + num_inference_steps=num_inference_steps, + guidance_scale=guidance_scale, + shift=shift, + timesteps_list=timesteps_list, + scheduler=scheduler, + generator=generator, + seed=seed, + noise_scale_start=noise_scale_start, + noise_scale_end=noise_scale_end, + noise_clip_std=noise_clip_std, + keep_original_aspect=keep_original_aspect, + callback=callback, + output_type=output_type, + return_dict=return_dict, + **kwargs, + ) diff --git a/pipelines/hidream/hidream_o1_utils.py b/pipelines/hidream/hidream_o1_utils.py new file mode 100644 index 000000000..b02e5231a --- /dev/null +++ b/pipelines/hidream/hidream_o1_utils.py @@ -0,0 +1,369 @@ +import math +from typing import Optional + +import numpy as np +import PIL.Image +import torch + + +TIMESTEP_TOKEN_NUM = 1 +NOISE_SCALE = 8.0 +T_EPS = 0.001 +CONDITION_IMAGE_SIZE = 384 +PATCH_SIZE = 32 + +DEFAULT_TIMESTEPS = [ + 999, + 987, + 974, + 960, + 945, + 929, + 913, + 895, + 877, + 857, + 836, + 814, + 790, + 764, + 737, + 707, + 675, + 640, + 602, + 560, + 515, + 464, + 409, + 347, + 278, + 199, + 110, + 8, +] + +PREDEFINED_RESOLUTIONS = [ + (2048, 2048), + (2304, 1728), + (1728, 2304), + (2560, 1440), + (1440, 2560), + (2496, 1664), + (1664, 2496), + (3104, 1312), + (1312, 3104), + (2304, 1792), + (1792, 2304), +] + + +def _ensure_special_tokens(tokenizer): + if not hasattr(tokenizer, "boi_token"): + tokenizer.boi_token = "<|boi_token|>" + if not hasattr(tokenizer, "bor_token"): + tokenizer.bor_token = "<|bor_token|>" + if not hasattr(tokenizer, "eor_token"): + tokenizer.eor_token = "<|eor_token|>" + if not hasattr(tokenizer, "bot_token"): + tokenizer.bot_token = "<|bot_token|>" + if not hasattr(tokenizer, "tms_token"): + tokenizer.tms_token = "<|tms_token|>" + + +def _find_closest_resolution(width: int, height: int): + img_ratio = width / height + best_res = PREDEFINED_RESOLUTIONS[0] + min_diff = float("inf") + for w, h in PREDEFINED_RESOLUTIONS: + diff = abs((w / h) - img_ratio) + if diff < min_diff: + min_diff = diff + best_res = (w, h) + return best_res + + +def _resize_pilimage( + pil_image: PIL.Image.Image, + image_size: int, + patch_size: int = 16, + resampler: PIL.Image.Resampling = PIL.Image.Resampling.BICUBIC, +): + while min(*pil_image.size) >= 2 * image_size: + pil_image = pil_image.resize(tuple(x // 2 for x in pil_image.size), resample=PIL.Image.Resampling.BOX) + + width, height = pil_image.width, pil_image.height + max_area = image_size * image_size + scale = math.sqrt(max_area / (width * height)) + + m = patch_size + new_sizes = [ + (round(width * scale) // m * m, round(height * scale) // m * m), + (round(width * scale) // m * m, math.floor(height * scale) // m * m), + (math.floor(width * scale) // m * m, round(height * scale) // m * m), + (math.floor(width * scale) // m * m, math.floor(height * scale) // m * m), + ] + new_sizes = sorted(new_sizes, key=lambda x: x[0] * x[1], reverse=True) + + new_size = new_sizes[-1] + for candidate in new_sizes: + if candidate[0] * candidate[1] <= max_area: + new_size = candidate + break + + s1 = width / new_size[0] + s2 = height / new_size[1] + if s1 < s2: + pil_image = pil_image.resize([new_size[0], round(height / s1)], resample=resampler) + top = (round(height / s1) - new_size[1]) // 2 + pil_image = pil_image.crop((0, top, new_size[0], top + new_size[1])) + else: + pil_image = pil_image.resize([round(width / s2), new_size[1]], resample=resampler) + left = (round(width / s2) - new_size[0]) // 2 + pil_image = pil_image.crop((left, 0, left + new_size[0], new_size[1])) + + return pil_image + + +def _calculate_dimensions(max_size: int, ratio: float): + width = math.sqrt(max_size * max_size * ratio) + height = width / ratio + width = int(width / 32) * 32 + height = int(height / 32) * 32 + return width, height + + +def _image_to_patch_tensor(x: torch.Tensor, patch_size: int = PATCH_SIZE) -> torch.Tensor: + b, c, h, w = x.shape + h_patch = h // patch_size + w_patch = w // patch_size + x = x.reshape(b, c, h_patch, patch_size, w_patch, patch_size) + x = x.permute(0, 2, 4, 1, 3, 5) + return x.reshape(b, h_patch * w_patch, c * patch_size * patch_size) + + +def _patch_tensor_to_image(x: torch.Tensor, h_patches: int, w_patches: int, patch_size: int = PATCH_SIZE) -> torch.Tensor: + b = x.shape[0] + c = x.shape[-1] // (patch_size * patch_size) + x = x.reshape(b, h_patches, w_patches, c, patch_size, patch_size) + x = x.permute(0, 3, 1, 4, 2, 5) + return x.reshape(b, c, h_patches * patch_size, w_patches * patch_size) + + +def _pil_to_normalized_tensor(image: PIL.Image.Image) -> torch.Tensor: + arr = np.array(image.convert("RGB"), dtype=np.float32) / 255.0 + x = torch.from_numpy(arr).permute(2, 0, 1) + return x * 2.0 - 1.0 + + +def _patches_to_np( + z: torch.Tensor, + h_patches: int, + w_patches: int, + invert: bool = True, + rescale: bool = False, +) -> np.ndarray: + z = z.float() + if rescale: + clip_ratio = (torch.abs(z) > 1.0).float().mean().item() # hidream-o1 often has outlines and clipping + if clip_ratio > 0.10: # Balanced quantiles and headroom: compress outliers while preserving contrast. + lo_q, hi_q, headroom = 0.03, 0.97, 1.33 + elif clip_ratio > 0.03: + lo_q, hi_q, headroom = 0.02, 0.98, 1.22 + else: + lo_q, hi_q, headroom = 0.01, 0.99, 1.15 + z_flat = z.reshape(-1) + q_lo = torch.quantile(z_flat, lo_q) + q_hi = torch.quantile(z_flat, hi_q) + center = (q_hi + q_lo) * 0.5 + half_range = (q_hi - q_lo) * 0.5 + if half_range > 0: + z = (z - center) / torch.clamp(half_range * headroom, min=0.35) # Balanced robust remap: moderate compression to handle outliers without losing contrast. + z = z.clamp(-1.0, 1.0) + image = (1.0 - z) / 2.0 if invert else (z + 1.0) / 2.0 + image = _patch_tensor_to_image(image, h_patches=h_patches, w_patches=w_patches, patch_size=PATCH_SIZE) + np_image = image[0].cpu().numpy().transpose(1, 2, 0) + np_image = np.round(255.0 * np_image).astype(np.uint8) + return np_image + + +def get_rope_index_fix_point( + spatial_merge_size, + image_token_id, + video_token_id, + vision_start_token_id, + 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, + skip_vision_start_token=None, + fix_point=4096, +) -> tuple[torch.Tensor, torch.Tensor]: + 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 + + 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_i in enumerate(total_input_ids): + input_ids_i = input_ids_i[attention_mask[i] == 1] + image_nums, video_nums = 0, 0 + vision_start_indices = torch.argwhere(input_ids_i == vision_start_token_id).squeeze(1) + vision_tokens = input_ids_i[vision_start_indices + 1] + image_nums = (vision_tokens == image_token_id).sum() + video_nums = (vision_tokens == video_token_id).sum() + input_tokens = input_ids_i.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] + image_index += 1 + remain_images -= 1 + ed = ed_image + else: + t, h, w = video_grid_thw[video_index] + 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 + + text_len -= 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 fix_point > 0: + fix_point = fix_point - st_idx + llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + fix_point + st_idx) + 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) + 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 + + 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 build_t2i_text_sample(prompt, height, width, tokenizer, processor, model_config): + 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) + + boi_token = getattr(tokenizer, "boi_token", "<|boi_token|>") + tms_token = getattr(tokenizer, "tms_token", "<|tms_token|>") + + messages = [{"role": "user", "content": prompt}] + template_caption = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + boi_token + tms_token * TIMESTEP_TOKEN_NUM + input_ids = tokenizer.encode(template_caption, return_tensors="pt", add_special_tokens=False) + + 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) + 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_fix_point( + 1, + image_token_id, + video_token_id, + vision_start_token_id, + input_ids=input_ids_pad, + image_grid_thw=image_grid_thw, + video_grid_thw=None, + attention_mask=None, + 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) + + return { + "input_ids": input_ids_pad, + "position_ids": position_ids, + "token_types": token_types_bin, + "vinput_mask": vinput_mask, + } + + +__all__ = [ + "CONDITION_IMAGE_SIZE", + "DEFAULT_TIMESTEPS", + "NOISE_SCALE", + "PATCH_SIZE", + "PREDEFINED_RESOLUTIONS", + "TIMESTEP_TOKEN_NUM", + "T_EPS", + "_calculate_dimensions", + "_ensure_special_tokens", + "_find_closest_resolution", + "_image_to_patch_tensor", + "_patch_tensor_to_image", + "_patches_to_np", + "_pil_to_normalized_tensor", + "_resize_pilimage", + "build_t2i_text_sample", + "get_rope_index_fix_point", +] diff --git a/pipelines/hidream/qwen3_vl_transformers.py b/pipelines/hidream/qwen3_vl_transformers.py new file mode 100644 index 000000000..d90d0cf41 --- /dev/null +++ b/pipelines/hidream/qwen3_vl_transformers.py @@ -0,0 +1,2065 @@ +# pylint: disable=unused-argument, protected-access +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") + +# Flash Attention import (FA3 preferred, FA2 fallback) +flash_attn_version = os.environ.get("FA_VERSION", "auto") +_flash_attn_func = None +if flash_attn_version == "2": + from flash_attn import flash_attn_func as _flash_attn_func +elif flash_attn_version == "3": + from flash_attn_interface import flash_attn_func as _flash_attn_func +else: + try: + from flash_attn_interface import flash_attn_func as _flash_attn_func + except ImportError: + try: + from flash_attn import flash_attn_func as _flash_attn_func + except ImportError: + _flash_attn_func = None + +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 check_model_inputs +from transformers.models.qwen3_vl.configuration_qwen3_vl import Qwen3VLConfig, Qwen3VLTextConfig, Qwen3VLVisionConfig + + +def _compute_default_rope_parameters(config, device=None, seq_len=None, layer_type=None): + rope_theta = getattr(config, "rope_theta", 10000.0) + partial_rotary_factor = getattr(config, "partial_rotary_factor", 1.0) + + if hasattr(config, "rope_parameters") and config.rope_parameters is not None: + if hasattr(config, "standardize_rope_params"): + config.standardize_rope_params() + rope_parameters = config.rope_parameters[layer_type] if layer_type is not None else config.rope_parameters + rope_theta = rope_parameters.get("rope_theta", rope_theta) + partial_rotary_factor = rope_parameters.get("partial_rotary_factor", partial_rotary_factor) + + head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads + dim = int(head_dim * partial_rotary_factor) + inv_freq = 1.0 / (rope_theta ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim)) + return inv_freq, 1.0 + +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 + rope_init_functions = dict(ROPE_INIT_FUNCTIONS) + rope_init_functions.setdefault("default", _compute_default_rope_parameters) + self.rope_init_fn = rope_init_functions.get(self.rope_type, _compute_default_rope_parameters) + + inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device) + self.register_buffer("inv_freq", inv_freq, persistent=False) + self.original_inv_freq = self.inv_freq + # self.original_inv_freq, _ = self.rope_init_fn(self.config, device) + + rope_scaling = config.rope_scaling if getattr(config, "rope_scaling", None) is not None else {} + self.mrope_section = rope_scaling.get("mrope_section", [24, 20, 20]) + + def compute_default_rope_parameters(self, config=None, device=None, seq_len=None, **kwargs): + rope_config = config or self.config + return _compute_default_rope_parameters(rope_config, device=device, seq_len=seq_len) + + def _materialize_inv_freq(self, device, use_original=False): + source = self.original_inv_freq if use_original else self.inv_freq + if source is not None and not source.is_meta: + return source.to(device=device) + + inv_freq, attention_scaling = self.rope_init_fn(self.config, device) + self.attention_scaling = attention_scaling + self.inv_freq = inv_freq + self.original_inv_freq = inv_freq.detach().clone() + return self.original_inv_freq if use_original else self.inv_freq + + 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 = self._materialize_inv_freq(x.device, use_original=False) + inv_freq_expanded = inv_freq[None, None, :, None].float().expand(3, position_ids.shape[1], -1, 1) + else: + inv_freq = self._materialize_inv_freq(x.device, use_original=True) + inv_freq_expanded = inv_freq[None, None, :, None].float().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() + + @check_model_inputs + @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 `(sequence_length)`, *optional*): + Cache positions used for KV-cache aware decoding and causal mask construction. + 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`, *optional*): + Decoder layer indices whose hidden states should be collected and returned in `mid_results`. + """ + 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, + inputs_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) === + _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 + + # 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 + + 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 flash attention two-pass approach. + + 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. + + 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 + """ + assert _flash_attn_func is not None, ( + "Flash attention is not available. Install flash_attn_interface (FA3) or flash_attn (FA2).") + + 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 _flash_layer_forward(hidden_states, decoder_layer, cos, sin, idx_ar): + """Flash 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 flash attention implementation. + """ + original_attn_forward = decoder_layer.self_attn.forward + + def _custom_flash_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 + q = attn.q_norm(attn.q_proj(hidden_states).view(hidden_shape)) + k = attn.k_norm(attn.k_proj(hidden_states).view(hidden_shape)) + v = attn.v_proj(hidden_states).view(hidden_shape) + + # Apply rotary position embedding (expects [B, H, S, D]) + cos_pe, sin_pe = position_embeddings + q_r = q.transpose(1, 2) # [B, H, S, D] + k_r = k.transpose(1, 2) # [B, KVH, S, D] + q_r, k_r = apply_rotary_pos_emb(q_r, k_r, cos_pe, sin_pe) + q = q_r.transpose(1, 2).contiguous() # [B, S, H, D] + k = k_r.transpose(1, 2).contiguous() # [B, S, KVH, D] + v = v.contiguous() + + softmax_scale = head_dim ** -0.5 + + # --- Two-pass flash attention --- + # Pass 1: causal attention on AR tokens only + q_ar = q[:, idx_ar].contiguous() + k_ar = k[:, idx_ar].contiguous() + v_ar = v[:, idx_ar].contiguous() + result_ar = _flash_attn_func(q_ar.to(torch.bfloat16), k_ar.to(torch.bfloat16), v_ar.to(torch.bfloat16), softmax_scale=softmax_scale, causal=True) + out_ar = result_ar[0] if isinstance(result_ar, tuple) else result_ar + + # Pass 2: full (bidirectional) attention on all tokens + result_full = _flash_attn_func(q.to(torch.bfloat16), k.to(torch.bfloat16), v.to(torch.bfloat16), softmax_scale=softmax_scale, causal=False) + out_full = result_full[0] if isinstance(result_full, tuple) else result_full + + # 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_flash_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( + _flash_layer_forward, + hidden_states, decoder_layer, cos, sin, idx_ar, + use_reentrant=False, + ) + else: + hidden_states = _flash_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] + batch_size, total_seq_len, _ = inputs_embeds.shape + + # Parse token_types early so we can distinguish conditioning vs generation placeholders. + parsed_token_types = None + if token_types is not None: + 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) + parsed_token_types = token_types + + # 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) + if parsed_token_types is not None: + conditioning_image_positions = (input_ids == self.config.image_token_id) & (~parsed_token_types.bool()) + conditioning_image_mask = conditioning_image_positions.unsqueeze(-1).expand_as(inputs_embeds) + if inputs_embeds[conditioning_image_mask].numel() != image_embeds.numel(): + raise ValueError( + "Image features and conditioning image tokens do not match: " + f"tokens: {conditioning_image_positions.sum()}, features {image_embeds.shape[0]}" + ) + inputs_embeds = inputs_embeds.masked_scatter(conditioning_image_mask, image_embeds) + else: + 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) + if parsed_token_types is not None: + conditioning_video_positions = (input_ids == self.config.video_token_id) & (~parsed_token_types.bool()) + conditioning_video_mask = conditioning_video_positions.unsqueeze(-1).expand_as(inputs_embeds) + if inputs_embeds[conditioning_video_mask].numel() != video_embeds.numel(): + raise ValueError( + "Video features and conditioning video tokens do not match: " + f"tokens: {conditioning_video_positions.sum()}, features {video_embeds.shape[0]}" + ) + inputs_embeds = inputs_embeds.masked_scatter(conditioning_video_mask, video_embeds) + else: + _, 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 fill the existing multimodal placeholder slots + 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] + + # 5. Parse token_types to [batch, total_seq_len] + token_types = parsed_token_types + if token_types is None: + raise ValueError("token_types is required for generation path") + + if vinputs_embedded.shape[0] == 1 and batch_size > 1: + vinputs_embedded = vinputs_embedded.expand(batch_size, -1, -1) + + vinput_positions = token_types.bool() & ~tms_mask + expected_vinput_tokens = vinputs_embedded.shape[1] + actual_vinput_tokens = vinput_positions.sum(dim=1) + if not torch.all(actual_vinput_tokens == expected_vinput_tokens): + raise ValueError( + f"Vinput token count mismatch: expected {expected_vinput_tokens} placeholder tokens, got {actual_vinput_tokens.tolist()}" + ) + + for batch_index in range(batch_size): + inputs_embeds[batch_index, vinput_positions[batch_index]] = vinputs_embedded[batch_index] + + # 6. Forward through decoder + mid_results = None + + 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] + + 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 + @check_model_inputs + 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""" + cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*): + Cache positions used during prefilling and incremental decoding. + 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. + vinputs (`torch.Tensor`, *optional*): + Patchified visual generation inputs appended to the token stream during denoising. + timestep (`torch.Tensor`, *optional*): + Per-sample denoising timesteps used to build timestep embeddings for generation. + token_types (`torch.Tensor`, *optional*): + Token-type mask distinguishing autoregressive text tokens from generation tokens. + use_flash_attn (`bool`, *optional*, defaults to `False`): + Whether to use the custom flash-attention generation path. + return_mid_results_layers (`list`, *optional*): + Decoder layer indices whose hidden states should be collected and returned in `mid_results`. + """ + # 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 HiDreamO1Qwen3VLTransformer(Qwen3VLPreTrainedModel, GenerationMixin): + _checkpoint_conversion_mapping = {} + _tied_weights_keys = ["lm_head.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 + + @check_model_inputs + 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/pipelines/hidream/scheduler_flashfloweuler.py b/pipelines/hidream/scheduler_flashfloweuler.py new file mode 100644 index 000000000..3d6f32c27 --- /dev/null +++ b/pipelines/hidream/scheduler_flashfloweuler.py @@ -0,0 +1,445 @@ +# Copyright 2024 Stability AI, Katherine Crowson and The HuggingFace 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 math +from dataclasses import dataclass +from typing import List, Optional, Tuple, Union + +import numpy as np +import torch +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.schedulers.scheduling_utils import SchedulerMixin +from diffusers.utils import BaseOutput, is_scipy_available, logging +from diffusers.utils.torch_utils import randn_tensor + +if is_scipy_available(): + import scipy.stats + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +def hack_randn_like(*args, **kwargs): + # for dist broadcast + value = torch.randn_like(*args, **kwargs) + return value + + +def hack_randn_tensor(*args, **kwargs): + # for dist broadcast + value = randn_tensor(*args, **kwargs) + return value + + +@dataclass +class FlashFlowMatchEulerDiscreteSchedulerOutput(BaseOutput): + """ + Output class for the scheduler's `step` function output. + + Args: + prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images): + Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the + denoising loop. + """ + + prev_sample: torch.FloatTensor + + +class FlashFlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin): + """ + Euler scheduler. + + This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic + methods the library implements for all schedulers such as loading and saving. + + Args: + num_train_timesteps (`int`, defaults to 1000): + The number of diffusion steps to train the model. + timestep_spacing (`str`, defaults to `"linspace"`): + The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and + Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information. + shift (`float`, defaults to 1.0): + The shift value for the timestep schedule. + """ + + _compatibles = [] + order = 1 + + @register_to_config + def __init__( + self, + num_train_timesteps: int = 1000, + shift: float = 1.0, + use_dynamic_shifting=True, + base_shift: Optional[float] = 0.5, + max_shift: Optional[float] = 1.15, + base_image_seq_len: Optional[int] = 256, + max_image_seq_len: Optional[int] = 4096, + invert_sigmas: bool = False, + use_karras_sigmas: Optional[bool] = False, + use_exponential_sigmas: Optional[bool] = False, + use_beta_sigmas: Optional[bool] = False, + ): + if self.config.use_beta_sigmas and not is_scipy_available(): + raise ImportError("Make sure to install scipy if you want to use beta sigmas.") + if sum([self.config.use_beta_sigmas, self.config.use_exponential_sigmas, self.config.use_karras_sigmas]) > 1: + raise ValueError( + "Only one of `config.use_beta_sigmas`, `config.use_exponential_sigmas`, `config.use_karras_sigmas` can be used." + ) + timesteps = np.linspace(1, num_train_timesteps, num_train_timesteps, dtype=np.float32)[::-1].copy() + timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32) + + sigmas = timesteps / num_train_timesteps + if not use_dynamic_shifting: + # when use_dynamic_shifting is True, we apply the timestep shifting on the fly based on the image resolution + sigmas = shift * sigmas / (1 + (shift - 1) * sigmas) + + self.timesteps = sigmas * num_train_timesteps + + self._step_index = None + self._begin_index = None + + self.sigmas = sigmas.to("cpu") # to avoid too much CPU/GPU communication + self.sigma_min = self.sigmas[-1].item() + self.sigma_max = self.sigmas[0].item() + + @property + def step_index(self): + """ + The index counter for current timestep. It will increase 1 after each scheduler step. + """ + return self._step_index + + @property + def begin_index(self): + """ + The index for the first timestep. It should be set from pipeline with `set_begin_index` method. + """ + return self._begin_index + + # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index + def set_begin_index(self, begin_index: int = 0): + """ + Sets the begin index for the scheduler. This function should be run from pipeline before the inference. + + Args: + begin_index (`int`): + The begin index for the scheduler. + """ + self._begin_index = begin_index + + def scale_noise( + self, + sample: torch.FloatTensor, + timestep: Union[float, torch.FloatTensor], + noise: Optional[torch.FloatTensor] = None, + ) -> torch.FloatTensor: + """ + Forward process in flow-matching + + Args: + sample (`torch.FloatTensor`): + The input sample. + timestep (`int`, *optional*): + The current timestep in the diffusion chain. + + Returns: + `torch.FloatTensor`: + A scaled input sample. + """ + # Make sure sigmas and timesteps have the same device and dtype as original_samples + sigmas = self.sigmas.to(device=sample.device, dtype=sample.dtype) + + if sample.device.type == "mps" and torch.is_floating_point(timestep): + # mps does not support float64 + schedule_timesteps = self.timesteps.to(sample.device, dtype=torch.float32) + timestep = timestep.to(sample.device, dtype=torch.float32) + else: + schedule_timesteps = self.timesteps.to(sample.device) + timestep = timestep.to(sample.device) + + # self.begin_index is None when scheduler is used for training, or pipeline does not implement set_begin_index + if self.begin_index is None: + step_indices = [self.index_for_timestep(t, schedule_timesteps) for t in timestep] + elif self.step_index is not None: + # add_noise is called after first denoising step (for inpainting) + step_indices = [self.step_index] * timestep.shape[0] + else: + # add noise is called before first denoising step to create initial latent(img2img) + step_indices = [self.begin_index] * timestep.shape[0] + + sigma = sigmas[step_indices].flatten() + while len(sigma.shape) < len(sample.shape): + sigma = sigma.unsqueeze(-1) + + sample = sigma * noise + (1.0 - sigma) * sample + + return sample + + def _sigma_to_t(self, sigma): + return sigma * self.config.num_train_timesteps + + def time_shift(self, mu: float, sigma: float, t: torch.Tensor): + return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma) + + def set_timesteps( + self, + num_inference_steps: int = None, + device: Union[str, torch.device] = None, + sigmas: Optional[List[float]] = None, + mu: Optional[float] = None, + ): + """ + Sets the discrete timesteps used for the diffusion chain (to be run before inference). + + Args: + num_inference_steps (`int`): + The number of diffusion steps used when generating samples with a pre-trained model. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + """ + if self.config.use_dynamic_shifting and mu is None: + raise ValueError(" you have a pass a value for `mu` when `use_dynamic_shifting` is set to be `True`") + + if sigmas is None: + timesteps = np.linspace( + self._sigma_to_t(self.sigma_max), self._sigma_to_t(self.sigma_min), num_inference_steps + ) + + sigmas = timesteps / self.config.num_train_timesteps + else: + sigmas = np.array(sigmas).astype(np.float32) + num_inference_steps = len(sigmas) + self.num_inference_steps = num_inference_steps + + if self.config.use_dynamic_shifting: + sigmas = self.time_shift(mu, 1.0, sigmas) + else: + sigmas = self.config.shift * sigmas / (1 + (self.config.shift - 1) * sigmas) + + if self.config.use_karras_sigmas: + sigmas = self._convert_to_karras(in_sigmas=sigmas, num_inference_steps=num_inference_steps) + + elif self.config.use_exponential_sigmas: + sigmas = self._convert_to_exponential(in_sigmas=sigmas, num_inference_steps=num_inference_steps) + + elif self.config.use_beta_sigmas: + sigmas = self._convert_to_beta(in_sigmas=sigmas, num_inference_steps=num_inference_steps) + + sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32, device=device) + timesteps = sigmas * self.config.num_train_timesteps + + if self.config.invert_sigmas: + sigmas = 1.0 - sigmas + timesteps = sigmas * self.config.num_train_timesteps + sigmas = torch.cat([sigmas, torch.ones(1, device=sigmas.device)]) + else: + sigmas = torch.cat([sigmas, torch.zeros(1, device=sigmas.device)]) + + self.timesteps = timesteps.to(device=device) + self.sigmas = sigmas + self._step_index = None + self._begin_index = None + + def index_for_timestep(self, timestep, schedule_timesteps=None): + if schedule_timesteps is None: + schedule_timesteps = self.timesteps + + indices = (schedule_timesteps == timestep).nonzero() + + # The sigma index that is taken for the **very** first `step` + # is always the second index (or the last index if there is only 1) + # This way we can ensure we don't accidentally skip a sigma in + # case we start in the middle of the denoising schedule (e.g. for image-to-image) + pos = 1 if len(indices) > 1 else 0 + + return indices[pos].item() + + def _init_step_index(self, timestep): + if self.begin_index is None: + if isinstance(timestep, torch.Tensor): + timestep = timestep.to(self.timesteps.device) + self._step_index = self.index_for_timestep(timestep) + else: + self._step_index = self._begin_index + + def step( + self, + model_output: torch.FloatTensor, + timestep: Union[float, torch.FloatTensor], + sample: torch.FloatTensor, + s_churn: float = 0.0, + s_tmin: float = 0.0, + s_tmax: float = float("inf"), + s_noise: float = 1.0, + noise_clip_std: float = 0.0, + generator: Optional[torch.Generator] = None, + return_dict: bool = True, + ) -> Union[FlashFlowMatchEulerDiscreteSchedulerOutput, Tuple]: + """ + Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion + process from the learned model outputs (most often the predicted noise). + + Args: + model_output (`torch.FloatTensor`): + The direct output from learned diffusion model. + timestep (`float`): + The current discrete timestep in the diffusion chain. + sample (`torch.FloatTensor`): + A current instance of a sample created by the diffusion process. + s_churn (`float`): + s_tmin (`float`): + s_tmax (`float`): + s_noise (`float`, defaults to 1.0): + Scaling factor for noise added to the sample. + generator (`torch.Generator`, *optional*): + A random number generator. + return_dict (`bool`): + Whether or not to return a [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or + tuple. + + Returns: + [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or `tuple`: + If return_dict is `True`, [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] is + returned, otherwise a tuple is returned where the first element is the sample tensor. + """ + + if ( + isinstance(timestep, int) + or isinstance(timestep, torch.IntTensor) + or isinstance(timestep, torch.LongTensor) + ): + raise ValueError( + ( + "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to" + " `EulerDiscreteScheduler.step()` is not supported. Make sure to pass" + " one of the `scheduler.timesteps` as a timestep." + ), + ) + + if self.step_index is None: + self._init_step_index(timestep) + + # Upcast to avoid precision issues when computing prev_sample + + sigma = self.sigmas[self.step_index] + + # Upcast to avoid precision issues when computing prev_sample + sample = sample.to(torch.float32) + + denoised = sample - model_output * sigma + + if self.step_index < self.num_inference_steps: + sigma_next = self.sigmas[self.step_index + 1] + noise = hack_randn_tensor( + model_output.shape, + generator=generator, + device=model_output.device, + dtype=denoised.dtype, + ) + if noise_clip_std > 0: + noise_std = noise.std().item() + clip_val = noise_clip_std * noise_std + noise = noise.clamp(min=-clip_val, max=clip_val) + sample = sigma_next * noise * s_noise + (1.0 - sigma_next) * denoised + + self._step_index += 1 + sample = sample.to(model_output.dtype) + + if not return_dict: + return (sample,) + + return FlashFlowMatchEulerDiscreteSchedulerOutput(prev_sample=sample) + + # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_karras + def _convert_to_karras(self, in_sigmas: torch.Tensor, num_inference_steps) -> torch.Tensor: + """Constructs the noise schedule of Karras et al. (2022).""" + + # Hack to make sure that other schedulers which copy this function don't break + # TODO: Add this logic to the other schedulers + if hasattr(self.config, "sigma_min"): + sigma_min = self.config.sigma_min + else: + sigma_min = None + + if hasattr(self.config, "sigma_max"): + sigma_max = self.config.sigma_max + else: + sigma_max = None + + sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() + sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() + + rho = 7.0 # 7.0 is the value used in the paper + ramp = np.linspace(0, 1, num_inference_steps) + min_inv_rho = sigma_min ** (1 / rho) + max_inv_rho = sigma_max ** (1 / rho) + sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho + return sigmas + + # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_exponential + def _convert_to_exponential(self, in_sigmas: torch.Tensor, num_inference_steps: int) -> torch.Tensor: + """Constructs an exponential noise schedule.""" + + # Hack to make sure that other schedulers which copy this function don't break + # TODO: Add this logic to the other schedulers + if hasattr(self.config, "sigma_min"): + sigma_min = self.config.sigma_min + else: + sigma_min = None + + if hasattr(self.config, "sigma_max"): + sigma_max = self.config.sigma_max + else: + sigma_max = None + + sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() + sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() + + sigmas = np.exp(np.linspace(math.log(sigma_max), math.log(sigma_min), num_inference_steps)) + return sigmas + + # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_beta + def _convert_to_beta( + self, in_sigmas: torch.Tensor, num_inference_steps: int, alpha: float = 0.6, beta: float = 0.6 + ) -> torch.Tensor: + """From "Beta Sampling is All You Need" [arXiv:2407.12173] (Lee et. al, 2024)""" + + # Hack to make sure that other schedulers which copy this function don't break + # TODO: Add this logic to the other schedulers + if hasattr(self.config, "sigma_min"): + sigma_min = self.config.sigma_min + else: + sigma_min = None + + if hasattr(self.config, "sigma_max"): + sigma_max = self.config.sigma_max + else: + sigma_max = None + + sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() + sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() + + sigmas = np.array( + [ + sigma_min + (ppf * (sigma_max - sigma_min)) + for ppf in [ + scipy.stats.beta.ppf(timestep, alpha, beta) + for timestep in 1 - np.linspace(0, 1, num_inference_steps) + ] + ] + ) + return sigmas + + def __len__(self): + return self.config.num_train_timesteps diff --git a/pipelines/hidream/scheduler_flowunipc.py b/pipelines/hidream/scheduler_flowunipc.py new file mode 100644 index 000000000..3b3d8ce29 --- /dev/null +++ b/pipelines/hidream/scheduler_flowunipc.py @@ -0,0 +1,778 @@ +# Copied from https://github.com/huggingface/diffusers/blob/v0.31.0/src/diffusers/schedulers/scheduling_unipc_multistep.py +# Convert unipc for flow matching +# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. + +import math +from typing import List, Optional, Tuple, Union + +import numpy as np +import torch +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.schedulers.scheduling_utils import (KarrasDiffusionSchedulers, + SchedulerMixin, + SchedulerOutput) + + +class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin): + """ + `UniPCMultistepScheduler` is a training-free framework designed for the fast sampling of diffusion models. + + This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic + methods the library implements for all schedulers such as loading and saving. + + Args: + num_train_timesteps (`int`, defaults to 1000): + The number of diffusion steps to train the model. + solver_order (`int`, default `2`): + The UniPC order which can be any positive integer. The effective order of accuracy is `solver_order + 1` + due to the UniC. It is recommended to use `solver_order=2` for guided sampling, and `solver_order=3` for + unconditional sampling. + prediction_type (`str`, defaults to "flow_prediction"): + Prediction type of the scheduler function; must be `flow_prediction` for this scheduler, which predicts + the flow of the diffusion process. + thresholding (`bool`, defaults to `False`): + Whether to use the "dynamic thresholding" method. This is unsuitable for latent-space diffusion models such + as Stable Diffusion. + dynamic_thresholding_ratio (`float`, defaults to 0.995): + The ratio for the dynamic thresholding method. Valid only when `thresholding=True`. + sample_max_value (`float`, defaults to 1.0): + The threshold value for dynamic thresholding. Valid only when `thresholding=True` and `predict_x0=True`. + predict_x0 (`bool`, defaults to `True`): + Whether to use the updating algorithm on the predicted x0. + solver_type (`str`, default `bh2`): + Solver type for UniPC. It is recommended to use `bh1` for unconditional sampling when steps < 10, and `bh2` + otherwise. + lower_order_final (`bool`, default `True`): + Whether to use lower-order solvers in the final steps. Only valid for < 15 inference steps. This can + stabilize the sampling of DPMSolver for steps < 15, especially for steps <= 10. + disable_corrector (`list`, default `[]`): + Decides which step to disable the corrector to mitigate the misalignment between `epsilon_theta(x_t, c)` + and `epsilon_theta(x_t^c, c)` which can influence convergence for a large guidance scale. Corrector is + usually disabled during the first few steps. + solver_p (`SchedulerMixin`, default `None`): + Any other scheduler that if specified, the algorithm becomes `solver_p + UniC`. + use_karras_sigmas (`bool`, *optional*, defaults to `False`): + Whether to use Karras sigmas for step sizes in the noise schedule during the sampling process. If `True`, + the sigmas are determined according to a sequence of noise levels {σi}. + use_exponential_sigmas (`bool`, *optional*, defaults to `False`): + Whether to use exponential sigmas for step sizes in the noise schedule during the sampling process. + timestep_spacing (`str`, defaults to `"linspace"`): + The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and + Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information. + steps_offset (`int`, defaults to 0): + An offset added to the inference steps, as required by some model families. + final_sigmas_type (`str`, defaults to `"zero"`): + The final `sigma` value for the noise schedule during the sampling process. If `"sigma_min"`, the final + sigma is the same as the last sigma in the training schedule. If `zero`, the final sigma is set to 0. + """ + + _compatibles = [e.name for e in KarrasDiffusionSchedulers] + order = 1 + + @register_to_config + def __init__( + self, + num_train_timesteps: int = 1000, + solver_order: int = 2, + prediction_type: str = "flow_prediction", + shift: Optional[float] = 1.0, + use_dynamic_shifting=False, + thresholding: bool = False, + dynamic_thresholding_ratio: float = 0.995, + sample_max_value: float = 1.0, + predict_x0: bool = True, + solver_type: str = "bh2", + lower_order_final: bool = True, + disable_corrector: List[int] = [], + solver_p: SchedulerMixin = None, + timestep_spacing: str = "linspace", + steps_offset: int = 0, + final_sigmas_type: Optional[str] = "zero", # "zero", "sigma_min" + ): + + if solver_type not in ["bh1", "bh2"]: + if solver_type in ["midpoint", "heun", "logrho"]: + self.register_to_config(solver_type="bh2") + else: + raise NotImplementedError( + f"{solver_type} is not implemented for {self.__class__}") + + self.predict_x0 = predict_x0 + # setable values + self.num_inference_steps = None + alphas = np.linspace(1, 1 / num_train_timesteps, + num_train_timesteps)[::-1].copy() + sigmas = 1.0 - alphas + sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32) + + if not use_dynamic_shifting: + # when use_dynamic_shifting is True, we apply the timestep shifting on the fly based on the image resolution + sigmas = shift * sigmas / (1 + + (shift - 1) * sigmas) # pyright: ignore + + self.sigmas = sigmas + self.timesteps = sigmas * num_train_timesteps + + self.model_outputs = [None] * solver_order + self.timestep_list = [None] * solver_order + self.lower_order_nums = 0 + self.disable_corrector = disable_corrector + self.solver_p = solver_p + self.last_sample = None + self._step_index = None + self._begin_index = None + + self.sigmas = self.sigmas.to( + "cpu") # to avoid too much CPU/GPU communication + self.sigma_min = self.sigmas[-1].item() + self.sigma_max = self.sigmas[0].item() + + @property + def step_index(self): + """ + The index counter for current timestep. It will increase 1 after each scheduler step. + """ + return self._step_index + + @property + def begin_index(self): + """ + The index for the first timestep. It should be set from pipeline with `set_begin_index` method. + """ + return self._begin_index + + # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index + def set_begin_index(self, begin_index: int = 0): + """ + Sets the begin index for the scheduler. This function should be run from pipeline before the inference. + + Args: + begin_index (`int`): + The begin index for the scheduler. + """ + self._begin_index = begin_index + + # Modified from diffusers.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler.set_timesteps + def set_timesteps( + self, + num_inference_steps: Union[int, None] = None, + device: Union[str, torch.device] = None, + sigmas: Optional[List[float]] = None, + mu: Optional[Union[float, None]] = None, + shift: Optional[Union[float, None]] = None, + ): + """ + Sets the discrete timesteps used for the diffusion chain (to be run before inference). + Args: + num_inference_steps (`int`): + Total number of the spacing of the time steps. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + """ + + if self.config.use_dynamic_shifting and mu is None: + raise ValueError( + " you have to pass a value for `mu` when `use_dynamic_shifting` is set to be `True`" + ) + + if sigmas is None: + sigmas = np.linspace(self.sigma_max, self.sigma_min, + num_inference_steps + + 1).copy()[:-1] # pyright: ignore + + if self.config.use_dynamic_shifting: + sigmas = self.time_shift(mu, 1.0, sigmas) # pyright: ignore + else: + if shift is None: + shift = self.config.shift + sigmas = shift * sigmas / (1 + + (shift - 1) * sigmas) # pyright: ignore + + if self.config.final_sigmas_type == "sigma_min": + sigma_last = ((1 - self.alphas_cumprod[0]) / + self.alphas_cumprod[0])**0.5 + elif self.config.final_sigmas_type == "zero": + sigma_last = 0 + else: + raise ValueError( + f"`final_sigmas_type` must be one of 'zero', or 'sigma_min', but got {self.config.final_sigmas_type}" + ) + + timesteps = sigmas * self.config.num_train_timesteps + sigmas = np.concatenate([sigmas, [sigma_last] + ]).astype(np.float32) # pyright: ignore + + self.sigmas = torch.from_numpy(sigmas) + self.timesteps = torch.from_numpy(timesteps).to( + device=device, dtype=torch.int64) + + self.num_inference_steps = len(timesteps) + + self.model_outputs = [ + None, + ] * self.config.solver_order + self.lower_order_nums = 0 + self.last_sample = None + if self.solver_p: + self.solver_p.set_timesteps(self.num_inference_steps, device=device) + + # add an index counter for schedulers that allow duplicated timesteps + self._step_index = None + self._begin_index = None + self.sigmas = self.sigmas.to( + "cpu") # to avoid too much CPU/GPU communication + + # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample + def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor: + """ + "Dynamic thresholding: At each sampling step we set s to a certain percentile absolute pixel value in xt0 (the + prediction of x_0 at timestep t), and if s > 1, then we threshold xt0 to the range [-s, s] and then divide by + s. Dynamic thresholding pushes saturated pixels (those near -1 and 1) inwards, thereby actively preventing + pixels from saturation at each step. We find that dynamic thresholding results in significantly better + photorealism as well as better image-text alignment, especially when using very large guidance weights." + + https://arxiv.org/abs/2205.11487 + """ + dtype = sample.dtype + batch_size, channels, *remaining_dims = sample.shape + + if dtype not in (torch.float32, torch.float64): + sample = sample.float( + ) # upcast for quantile calculation, and clamp not implemented for cpu half + + # Flatten sample for doing quantile calculation along each image + sample = sample.reshape(batch_size, channels * np.prod(remaining_dims)) + + abs_sample = sample.abs() # "a certain percentile absolute pixel value" + + s = torch.quantile( + abs_sample, self.config.dynamic_thresholding_ratio, dim=1) + s = torch.clamp( + s, min=1, max=self.config.sample_max_value + ) # When clamped to min=1, equivalent to standard clipping to [-1, 1] + s = s.unsqueeze( + 1) # (batch_size, 1) because clamp will broadcast along dim=0 + sample = torch.clamp( + sample, -s, s + ) / s # "we threshold xt0 to the range [-s, s] and then divide by s" + + sample = sample.reshape(batch_size, channels, *remaining_dims) + sample = sample.to(dtype) + + return sample + + # Copied from diffusers.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler._sigma_to_t + def _sigma_to_t(self, sigma): + return sigma * self.config.num_train_timesteps + + def _sigma_to_alpha_sigma_t(self, sigma): + return 1 - sigma, sigma + + # Copied from diffusers.schedulers.scheduling_flow_match_euler_discrete.set_timesteps + def time_shift(self, mu: float, sigma: float, t: torch.Tensor): + return math.exp(mu) / (math.exp(mu) + (1 / t - 1)**sigma) + + def convert_model_output( + self, + model_output: torch.Tensor, + *args, + sample: torch.Tensor = None, + **kwargs, + ) -> torch.Tensor: + r""" + Convert the model output to the corresponding type the UniPC algorithm needs. + + Args: + model_output (`torch.Tensor`): + The direct output from the learned diffusion model. + timestep (`int`): + The current discrete timestep in the diffusion chain. + sample (`torch.Tensor`): + A current instance of a sample created by the diffusion process. + + Returns: + `torch.Tensor`: + The converted model output. + """ + timestep = args[0] if len(args) > 0 else kwargs.pop("timestep", None) + if sample is None: + if len(args) > 1: + sample = args[1] + else: + raise ValueError( + "missing `sample` as a required keyward argument") + + sigma = self.sigmas[self.step_index] + alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma) + + if self.predict_x0: + if self.config.prediction_type == "flow_prediction": + sigma_t = self.sigmas[self.step_index] + x0_pred = sample - sigma_t * model_output + else: + raise ValueError( + f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`," + " `v_prediction` or `flow_prediction` for the UniPCMultistepScheduler." + ) + + if self.config.thresholding: + x0_pred = self._threshold_sample(x0_pred) + + return x0_pred + else: + if self.config.prediction_type == "flow_prediction": + sigma_t = self.sigmas[self.step_index] + epsilon = sample - (1 - sigma_t) * model_output + else: + raise ValueError( + f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`," + " `v_prediction` or `flow_prediction` for the UniPCMultistepScheduler." + ) + + if self.config.thresholding: + sigma_t = self.sigmas[self.step_index] + x0_pred = sample - sigma_t * model_output + x0_pred = self._threshold_sample(x0_pred) + epsilon = model_output + x0_pred + + return epsilon + + def multistep_uni_p_bh_update( + self, + model_output: torch.Tensor, + *args, + sample: torch.Tensor = None, + order: int = None, # pyright: ignore + **kwargs, + ) -> torch.Tensor: + """ + One step for the UniP (B(h) version). Alternatively, `self.solver_p` is used if is specified. + + Args: + model_output (`torch.Tensor`): + The direct output from the learned diffusion model at the current timestep. + prev_timestep (`int`): + The previous discrete timestep in the diffusion chain. + sample (`torch.Tensor`): + A current instance of a sample created by the diffusion process. + order (`int`): + The order of UniP at this timestep (corresponds to the *p* in UniPC-p). + + Returns: + `torch.Tensor`: + The sample tensor at the previous timestep. + """ + prev_timestep = args[0] if len(args) > 0 else kwargs.pop( + "prev_timestep", None) + if sample is None: + if len(args) > 1: + sample = args[1] + else: + raise ValueError( + " missing `sample` as a required keyward argument") + if order is None: + if len(args) > 2: + order = args[2] + else: + raise ValueError( + " missing `order` as a required keyward argument") + model_output_list = self.model_outputs + + s0 = self.timestep_list[-1] + m0 = model_output_list[-1] + x = sample + + if self.solver_p: + x_t = self.solver_p.step(model_output, s0, x).prev_sample + return x_t + + sigma_t, sigma_s0 = self.sigmas[self.step_index + 1], self.sigmas[ + self.step_index] # pyright: ignore + alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t) + alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0) + + lambda_t = torch.log(alpha_t) - torch.log(sigma_t) + lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0) + + h = lambda_t - lambda_s0 + device = sample.device + + rks = [] + D1s = [] + for i in range(1, order): + si = self.step_index - i # pyright: ignore + mi = model_output_list[-(i + 1)] + alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si]) + lambda_si = torch.log(alpha_si) - torch.log(sigma_si) + rk = (lambda_si - lambda_s0) / h + rks.append(rk) + D1s.append((mi - m0) / rk) # pyright: ignore + + rks.append(1.0) + rks = torch.tensor(rks, device=device) + + R = [] + b = [] + + hh = -h if self.predict_x0 else h + h_phi_1 = torch.expm1(hh) # h\phi_1(h) = e^h - 1 + h_phi_k = h_phi_1 / hh - 1 + + factorial_i = 1 + + if self.config.solver_type == "bh1": + B_h = hh + elif self.config.solver_type == "bh2": + B_h = torch.expm1(hh) + else: + raise NotImplementedError + + for i in range(1, order + 1): + R.append(torch.pow(rks, i - 1)) + b.append(h_phi_k * factorial_i / B_h) + factorial_i *= i + 1 + h_phi_k = h_phi_k / hh - 1 / factorial_i + + R = torch.stack(R) + b = torch.tensor(b, device=device) + + if len(D1s) > 0: + D1s = torch.stack(D1s, dim=1) # (B, K) + # for order 2, we use a simplified version + if order == 2: + rhos_p = torch.tensor([0.5], dtype=x.dtype, device=device) + else: + rhos_p = torch.linalg.solve(R[:-1, :-1], + b[:-1]).to(device).to(x.dtype) + else: + D1s = None + + if self.predict_x0: + x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0 + if D1s is not None: + pred_res = torch.einsum("k,bkc...->bc...", rhos_p, + D1s) # pyright: ignore + else: + pred_res = 0 + x_t = x_t_ - alpha_t * B_h * pred_res + else: + x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0 + if D1s is not None: + pred_res = torch.einsum("k,bkc...->bc...", rhos_p, + D1s) # pyright: ignore + else: + pred_res = 0 + x_t = x_t_ - sigma_t * B_h * pred_res + + x_t = x_t.to(x.dtype) + return x_t + + def multistep_uni_c_bh_update( + self, + this_model_output: torch.Tensor, + *args, + last_sample: torch.Tensor = None, + this_sample: torch.Tensor = None, + order: int = None, # pyright: ignore + **kwargs, + ) -> torch.Tensor: + """ + One step for the UniC (B(h) version). + + Args: + this_model_output (`torch.Tensor`): + The model outputs at `x_t`. + this_timestep (`int`): + The current timestep `t`. + last_sample (`torch.Tensor`): + The generated sample before the last predictor `x_{t-1}`. + this_sample (`torch.Tensor`): + The generated sample after the last predictor `x_{t}`. + order (`int`): + The `p` of UniC-p at this step. The effective order of accuracy should be `order + 1`. + + Returns: + `torch.Tensor`: + The corrected sample tensor at the current timestep. + """ + this_timestep = args[0] if len(args) > 0 else kwargs.pop( + "this_timestep", None) + if last_sample is None: + if len(args) > 1: + last_sample = args[1] + else: + raise ValueError( + " missing`last_sample` as a required keyward argument") + if this_sample is None: + if len(args) > 2: + this_sample = args[2] + else: + raise ValueError( + " missing`this_sample` as a required keyward argument") + if order is None: + if len(args) > 3: + order = args[3] + else: + raise ValueError( + " missing`order` as a required keyward argument") + + model_output_list = self.model_outputs + + m0 = model_output_list[-1] + x = last_sample + x_t = this_sample + model_t = this_model_output + + sigma_t, sigma_s0 = self.sigmas[self.step_index], self.sigmas[ + self.step_index - 1] # pyright: ignore + alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t) + alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0) + + lambda_t = torch.log(alpha_t) - torch.log(sigma_t) + lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0) + + h = lambda_t - lambda_s0 + device = this_sample.device + + rks = [] + D1s = [] + for i in range(1, order): + si = self.step_index - (i + 1) # pyright: ignore + mi = model_output_list[-(i + 1)] + alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si]) + lambda_si = torch.log(alpha_si) - torch.log(sigma_si) + rk = (lambda_si - lambda_s0) / h + rks.append(rk) + D1s.append((mi - m0) / rk) # pyright: ignore + + rks.append(1.0) + rks = torch.tensor(rks, device=device) + + R = [] + b = [] + + hh = -h if self.predict_x0 else h + h_phi_1 = torch.expm1(hh) # h\phi_1(h) = e^h - 1 + h_phi_k = h_phi_1 / hh - 1 + + factorial_i = 1 + + if self.config.solver_type == "bh1": + B_h = hh + elif self.config.solver_type == "bh2": + B_h = torch.expm1(hh) + else: + raise NotImplementedError + + for i in range(1, order + 1): + R.append(torch.pow(rks, i - 1)) + b.append(h_phi_k * factorial_i / B_h) + factorial_i *= i + 1 + h_phi_k = h_phi_k / hh - 1 / factorial_i + + R = torch.stack(R) + b = torch.tensor(b, device=device) + + if len(D1s) > 0: + D1s = torch.stack(D1s, dim=1) + else: + D1s = None + + # for order 1, we use a simplified version + if order == 1: + rhos_c = torch.tensor([0.5], dtype=x.dtype, device=device) + else: + rhos_c = torch.linalg.solve(R, b).to(device).to(x.dtype) + + if self.predict_x0: + x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0 + if D1s is not None: + corr_res = torch.einsum("k,bkc...->bc...", rhos_c[:-1], D1s) + else: + corr_res = 0 + D1_t = model_t - m0 + x_t = x_t_ - alpha_t * B_h * (corr_res + rhos_c[-1] * D1_t) + else: + x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0 + if D1s is not None: + corr_res = torch.einsum("k,bkc...->bc...", rhos_c[:-1], D1s) + else: + corr_res = 0 + D1_t = model_t - m0 + x_t = x_t_ - sigma_t * B_h * (corr_res + rhos_c[-1] * D1_t) + x_t = x_t.to(x.dtype) + return x_t + + def index_for_timestep(self, timestep, schedule_timesteps=None): + if schedule_timesteps is None: + schedule_timesteps = self.timesteps + + indices = (schedule_timesteps == timestep).nonzero() + + # The sigma index that is taken for the **very** first `step` + # is always the second index (or the last index if there is only 1) + # This way we can ensure we don't accidentally skip a sigma in + # case we start in the middle of the denoising schedule (e.g. for image-to-image) + pos = 1 if len(indices) > 1 else 0 + + return indices[pos].item() + + # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler._init_step_index + def _init_step_index(self, timestep): + """ + Initialize the step_index counter for the scheduler. + """ + + if self.begin_index is None: + if isinstance(timestep, torch.Tensor): + timestep = timestep.to(self.timesteps.device) + self._step_index = self.index_for_timestep(timestep) + else: + self._step_index = self._begin_index + + def step(self, + model_output: torch.Tensor, + timestep: Union[int, torch.Tensor], + sample: torch.Tensor, + return_dict: bool = True, + generator=None) -> Union[SchedulerOutput, Tuple]: + """ + Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with + the multistep UniPC. + + Args: + model_output (`torch.Tensor`): + The direct output from learned diffusion model. + timestep (`int`): + The current discrete timestep in the diffusion chain. + sample (`torch.Tensor`): + A current instance of a sample created by the diffusion process. + return_dict (`bool`): + Whether or not to return a [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`. + + Returns: + [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`: + If return_dict is `True`, [`~schedulers.scheduling_utils.SchedulerOutput`] is returned, otherwise a + tuple is returned where the first element is the sample tensor. + + """ + if self.num_inference_steps is None: + raise ValueError( + "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler" + ) + + if self.step_index is None: + self._init_step_index(timestep) + + use_corrector = ( + self.step_index > 0 and + self.step_index - 1 not in self.disable_corrector and + self.last_sample is not None # pyright: ignore + ) + + model_output_convert = self.convert_model_output( + model_output, sample=sample) + if use_corrector: + sample = self.multistep_uni_c_bh_update( + this_model_output=model_output_convert, + last_sample=self.last_sample, + this_sample=sample, + order=self.this_order, + ) + + for i in range(self.config.solver_order - 1): + self.model_outputs[i] = self.model_outputs[i + 1] + self.timestep_list[i] = self.timestep_list[i + 1] + + self.model_outputs[-1] = model_output_convert + self.timestep_list[-1] = timestep # pyright: ignore + + if self.config.lower_order_final: + this_order = min(self.config.solver_order, + len(self.timesteps) - + self.step_index) # pyright: ignore + else: + this_order = self.config.solver_order + + self.this_order = min(this_order, + self.lower_order_nums + 1) # warmup for multistep + assert self.this_order > 0 + + self.last_sample = sample + prev_sample = self.multistep_uni_p_bh_update( + model_output=model_output, # pass the original non-converted model output, in case solver-p is used + sample=sample, + order=self.this_order, + ) + + if self.lower_order_nums < self.config.solver_order: + self.lower_order_nums += 1 + + # upon completion increase step index by one + self._step_index += 1 # pyright: ignore + + if not return_dict: + return (prev_sample,) + + return SchedulerOutput(prev_sample=prev_sample) + + def scale_model_input(self, sample: torch.Tensor, *args, + **kwargs) -> torch.Tensor: + """ + Ensures interchangeability with schedulers that need to scale the denoising model input depending on the + current timestep. + + Args: + sample (`torch.Tensor`): + The input sample. + + Returns: + `torch.Tensor`: + A scaled input sample. + """ + return sample + + # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.add_noise + def add_noise( + self, + original_samples: torch.Tensor, + noise: torch.Tensor, + timesteps: torch.IntTensor, + ) -> torch.Tensor: + # Make sure sigmas and timesteps have the same device and dtype as original_samples + sigmas = self.sigmas.to( + device=original_samples.device, dtype=original_samples.dtype) + if original_samples.device.type == "mps" and torch.is_floating_point( + timesteps): + # mps does not support float64 + schedule_timesteps = self.timesteps.to( + original_samples.device, dtype=torch.float32) + timesteps = timesteps.to( + original_samples.device, dtype=torch.float32) + else: + schedule_timesteps = self.timesteps.to(original_samples.device) + timesteps = timesteps.to(original_samples.device) + + # begin_index is None when the scheduler is used for training or pipeline does not implement set_begin_index + if self.begin_index is None: + step_indices = [ + self.index_for_timestep(t, schedule_timesteps) + for t in timesteps + ] + elif self.step_index is not None: + # add_noise is called after first denoising step (for inpainting) + step_indices = [self.step_index] * timesteps.shape[0] + else: + # add noise is called before first denoising step to create initial latent(img2img) + step_indices = [self.begin_index] * timesteps.shape[0] + + sigma = sigmas[step_indices].flatten() + while len(sigma.shape) < len(original_samples.shape): + sigma = sigma.unsqueeze(-1) + + alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma) + noisy_samples = alpha_t * original_samples + sigma_t * noise + return noisy_samples + + def __len__(self): + return self.config.num_train_timesteps diff --git a/pipelines/model_hidream.py b/pipelines/model_hidream.py index 1e05fbfef..45bacd81f 100644 --- a/pipelines/model_hidream.py +++ b/pipelines/model_hidream.py @@ -31,6 +31,56 @@ def load_llama(diffusers_load_config=None): return text_encoder_4, tokenizer_4 +def load_hidream_o1(checkpoint_info, diffusers_load_config=None): + if diffusers_load_config is None: + diffusers_load_config = {} + repo_id = sd_models.path_to_repo(checkpoint_info) + sd_models.hf_auth_check(checkpoint_info) + + from pipelines.hidream.hidream_o1 import HiDreamO1Pipeline, HiDreamO1ImagePipeline + from pipelines.hidream.qwen3_vl_transformers import HiDreamO1Qwen3VLTransformer + from pipelines.hidream.scheduler_flashfloweuler import FlashFlowMatchEulerDiscreteScheduler + + load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='Model', device_map=True, allow_quant=False) + log.debug(f'Load model: type=HiDreamO1 repo="{repo_id}" offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') + + o1_load_config = diffusers_load_config.copy() + o1_load_config['trust_remote_code'] = True + + transformer = HiDreamO1Qwen3VLTransformer.from_pretrained( + repo_id, + cache_dir=shared.opts.hfcache_dir, + trust_remote_code=True, + **load_args, + **quant_args, + ) + if shared.opts.diffusers_offload_mode != 'none' and transformer is not None: + sd_models.move_model(transformer, devices.cpu) + + processor = transformers.AutoProcessor.from_pretrained( + repo_id, + cache_dir=shared.opts.hfcache_dir, + trust_remote_code=True, + ) + pipe = HiDreamO1Pipeline( + transformer=transformer, + processor=processor, + tokenizer=processor.tokenizer, + scheduler=FlashFlowMatchEulerDiscreteScheduler(num_train_timesteps=1000, shift=3.0, use_dynamic_shifting=False), + ) + pipe.task_args = { + 'output_type': 'pil', + } + + del processor + del transformer + diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["hidream-o1"] = HiDreamO1Pipeline + diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["hidream-o1"] = HiDreamO1ImagePipeline + + devices.torch_gc() + return pipe + + def load_hidream(checkpoint_info, diffusers_load_config=None): if diffusers_load_config is None: diffusers_load_config = {} @@ -52,7 +102,7 @@ def load_hidream(checkpoint_info, diffusers_load_config=None): if 'I1' in repo_id: cls = diffusers.HiDreamImagePipeline elif 'E1' in repo_id: - from pipelines.hidream.pipeline_hidream_image_editing import HiDreamImageEditingPipeline + from pipelines.hidream.hidream_e1 import HiDreamImageEditingPipeline cls = HiDreamImageEditingPipeline diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["hidream-e1"] = diffusers.HiDreamImagePipeline diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["hidream-e1"] = HiDreamImageEditingPipeline diff --git a/pyproject.toml b/pyproject.toml index 1ecf3f25e..eb2fd67e8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -30,6 +30,7 @@ exclude = [ "pipelines/meissonic", "pipelines/omnigen2", "pipelines/hdm", + "pipelines/hidream", "pipelines/segmoe", "pipelines/xomni", "pipelines/chrono",