# Copyright 2026 Bytedance Ltd. and/or its affiliates # # 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. from __future__ import annotations import copy import os from typing import Any, Literal import torch from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.models.qwen_image import QwenImagePipeline from vllm_omni.diffusion.models.qwen_image.rope_utils import txt_seq_lens_from_embeds from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.diffusion.worker.utils import StepRequestState from verl_omni.pipelines.diffusion_rollout_output import rollout_output, with_rollout_data from verl_omni.pipelines.model_base import VllmOmniPipelineBase from verl_omni.pipelines.request_batch import ( collate_prompt_mask as _collate_prompt_mask, ) from verl_omni.pipelines.request_batch import ( collate_prompt_rows as _collate_prompt_rows, ) from verl_omni.pipelines.request_batch import ( sample_per_sample_sde_windows as _sample_per_sample_sde_windows, ) from verl_omni.pipelines.request_batch import ( split_diffusion_output_by_request as _split_diffusion_output_by_request, ) from verl_omni.pipelines.rollout_media import DiffusionIOSpec, MediaSpec from verl_omni.pipelines.schedulers import FlowMatchSDEDiscreteScheduler from .common import QwenImageTokenIdPromptMixin, apply_true_cfg, build_img_shapes, coalesce_not_none __all__ = ["QwenImagePipelineWithLogProb"] @VllmOmniPipelineBase.register("QwenImagePipeline", algorithm="flow_grpo") class QwenImagePipelineWithLogProb(QwenImageTokenIdPromptMixin, QwenImagePipeline): """Rollout pipeline for Qwen-Image that captures per-step log-probabilities. Extends :class:`~vllm_omni.diffusion.models.qwen_image.QwenImagePipeline` with a custom SDE-based scheduler and additional output fields required for RL training (e.g. FlowGRPO). In addition to the final generated image the pipeline returns all intermediate latents, their log-probabilities, and the corresponding timesteps. Registered under ``"QwenImagePipeline"`` for vllm-omni rollout dispatch. """ supports_request_batch = True #: Declares the primary rollout media stream so downstream consumers read #: the modality from the adapter instead of inferring it from tensor rank. #: Inherited by the dual-GRPO and mix-GRPO Qwen-Image subclasses. diffusion_io_spec = DiffusionIOSpec(primary=MediaSpec("image")) def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): super().__init__(od_config=od_config, prefix=prefix) self.device = get_local_device() model = od_config.model local_files_only = os.path.exists(model) self.scheduler = FlowMatchSDEDiscreteScheduler.from_pretrained( model, subfolder="scheduler", local_files_only=local_files_only, ) def _get_qwen_prompt_embeds( self, prompt_ids: torch.Tensor, attention_mask: torch.Tensor | None = None, dtype: torch.dtype | None = None, ): dtype = dtype or self.text_encoder.dtype if attention_mask is None: attention_mask = torch.ones_like(prompt_ids, dtype=torch.long) prompt_ids = prompt_ids.unsqueeze(0) if prompt_ids.ndim == 1 else prompt_ids attention_mask = attention_mask.unsqueeze(0) if attention_mask.ndim == 1 else attention_mask drop_idx = self.prompt_template_encode_start_idx encoder_hidden_states = self.text_encoder( input_ids=prompt_ids.to(self.device), attention_mask=attention_mask.to(self.device), output_hidden_states=True, ) hidden_states = encoder_hidden_states.hidden_states[-1] split_hidden_states = self._extract_masked_hidden(hidden_states, attention_mask) split_hidden_states = [e[drop_idx:] for e in split_hidden_states] attn_mask_list = [torch.ones(e.size(0), dtype=torch.long, device=e.device) for e in split_hidden_states] max_seq_len = max([e.size(0) for e in split_hidden_states]) prompt_embeds = torch.stack( [torch.cat([u, u.new_zeros(max_seq_len - u.size(0), u.size(1))]) for u in split_hidden_states] ) encoder_attention_mask = torch.stack( [torch.cat([u, u.new_zeros(max_seq_len - u.size(0))]) for u in attn_mask_list] ) prompt_embeds = prompt_embeds.to(dtype=dtype) return prompt_embeds, encoder_attention_mask def encode_prompt( self, prompt_ids: torch.Tensor | list[int] | list[list[int]] | None, attention_mask: torch.Tensor | list[int] | list[bool] | list[list[int]] | list[list[bool]] | None = None, num_images_per_prompt: int = 1, prompt_embeds: torch.Tensor | None = None, prompt_embeds_mask: torch.Tensor | None = None, max_sequence_length: int = 1024, ): """Encode text prompt token IDs into dense embeddings. Args: prompt_ids (torch.Tensor | list): Token IDs of shape ``(B, L)`` or ``(L,)``. attention_mask (torch.Tensor | list, *optional*): Boolean mask of shape ``(B, L)`` for *prompt_ids*; inferred as all-ones when ``None``. num_images_per_prompt (int): Number of images to generate per prompt; embeddings are repeated accordingly. prompt_embeds (torch.Tensor, *optional*): Pre-computed embeddings; when provided *prompt_ids* is ignored. prompt_embeds_mask (torch.Tensor, *optional*): Attention mask for pre-computed *prompt_embeds*. max_sequence_length (int): Maximum sequence length; embeddings are truncated to this value. Returns: tuple[torch.Tensor, torch.Tensor]: A pair of ``(prompt_embeds, prompt_embeds_mask)`` tensors of shape ``(B * num_images_per_prompt, L, D)`` and ``(B * num_images_per_prompt, L)`` respectively. """ if prompt_embeds is None: if prompt_ids is None: raise ValueError("`prompt_ids` must be provided when `prompt_embeds` is None.") if isinstance(prompt_ids, list): prompt_ids = torch.tensor(prompt_ids, device=self.device, dtype=torch.long) elif isinstance(prompt_ids, torch.Tensor): prompt_ids = prompt_ids.to(self.device) else: raise TypeError("`prompt_ids` must be a tensor or list.") if isinstance(attention_mask, list): attention_mask = torch.tensor(attention_mask, device=self.device) elif isinstance(attention_mask, torch.Tensor): attention_mask = attention_mask.to(self.device) elif attention_mask is not None: raise TypeError("`attention_mask` must be a tensor or list.") prompt_ids = prompt_ids.unsqueeze(0) if prompt_ids.ndim == 1 else prompt_ids attention_mask = ( attention_mask.unsqueeze(0) if attention_mask is not None and attention_mask.ndim == 1 else attention_mask ) prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds(prompt_ids, attention_mask=attention_mask) prompt_embeds = prompt_embeds[:, :max_sequence_length] prompt_embeds_mask = prompt_embeds_mask[:, :max_sequence_length] if num_images_per_prompt > 1: prompt_embeds = prompt_embeds.repeat_interleave(num_images_per_prompt, dim=0) prompt_embeds_mask = prompt_embeds_mask.repeat_interleave(num_images_per_prompt, dim=0) return prompt_embeds, prompt_embeds_mask def _extract_prompt_ids(self, prompts): """Extract prompt_ids/mask and their negatives from the OmniCustomPrompt list. Falls back to tokenizing ``"prompt"`` / ``"negative_prompt"`` text fields when ``prompt_ids`` is not provided (e.g. during the engine's dummy warm-up run, which always submits a text prompt). """ prompt_ids = None prompt_mask = None negative_prompt_ids = None negative_prompt_mask = None if prompts: p0 = prompts[0] if isinstance(p0, dict): prompt_ids = p0.get("prompt_token_ids", None) prompt_mask = p0.get("prompt_mask", None) negative_prompt_ids = p0.get("negative_prompt_ids", None) negative_prompt_mask = p0.get("negative_prompt_mask", None) # Fallback: tokenize raw text prompt (covers _dummy_run path). if prompt_ids is None and p0.get("prompt"): prompt_ids, prompt_mask = self._tokenize_text_prompt(p0["prompt"]) if negative_prompt_ids is None and p0.get("negative_prompt"): negative_prompt_ids, negative_prompt_mask = self._tokenize_text_prompt(p0["negative_prompt"]) elif isinstance(p0, str): prompt_ids, prompt_mask = self._tokenize_text_prompt(p0) return prompt_ids, prompt_mask, negative_prompt_ids, negative_prompt_mask def _tokenize_text_prompt(self, text: str | list[str]): """Tokenize a text prompt using the Qwen chat template (parent behavior).""" prompt = [text] if isinstance(text, str) else text txt = [self.prompt_template_encode.format(e) for e in prompt] prompt_embed_cache = getattr(self, "_prompt_embed_cache", None) use_prompt_embed_cache = bool(prompt_embed_cache is not None and prompt_embed_cache.enabled) tokenizer_kwargs = {} if use_prompt_embed_cache else {"return_tensors": "pt"} tokens = self.tokenizer( txt, max_length=self.tokenizer_max_length + self.prompt_template_encode_start_idx, padding=True, truncation=True, **tokenizer_kwargs, ) if not use_prompt_embed_cache: tokens = tokens.to(self.device) return tokens.input_ids, tokens.attention_mask def prepare_encode( self, state: StepRequestState, **kwargs: Any, ) -> StepRequestState: """Populate *state* with encoded prompts, latents, timesteps, and CFG config. Override of ``QwenImagePipeline.prepare_encode`` that accepts pre-tokenized ``prompt_ids`` (and optional ``prompt_mask``) instead of raw text prompts, matching the input contract of ``QwenImagePipelineWithLogProb``. """ sampling = state.sampling # vllm-omni >=0.24 stores a single prompt on StepRequestState.prompt # (not .prompts). Match upstream QwenImagePipeline.prepare_encode. prompt_ids, prompt_mask, negative_prompt_ids, negative_prompt_mask = self._extract_prompt_ids( [state.prompt] if state.prompt is not None else [] ) if prompt_ids is None: raise ValueError( f"{self.__class__.__name__}.prepare_encode requires either " "'prompt_ids'/'prompt_token_ids' or a text 'prompt' on state.prompt." ) height = sampling.height or self.default_sample_size * self.vae_scale_factor width = sampling.width or self.default_sample_size * self.vae_scale_factor num_inference_steps = sampling.num_inference_steps or 50 sigmas = sampling.sigmas guidance_scale = sampling.guidance_scale if sampling.guidance_scale_provided else 1.0 num_images_per_prompt = sampling.num_outputs_per_prompt if sampling.num_outputs_per_prompt > 0 else 1 true_cfg_scale = sampling.true_cfg_scale or 4.0 max_sequence_length = sampling.max_sequence_length or self.tokenizer_max_length generator = sampling.generator if generator is None and sampling.seed is not None: generator = torch.Generator(device=self.device).manual_seed(sampling.seed) self._guidance_scale = guidance_scale self._attention_kwargs = kwargs.get("attention_kwargs") or {} self._current_timestep = None self._interrupt = False if isinstance(prompt_ids, torch.Tensor): batch_size = prompt_ids.shape[0] if prompt_ids.ndim == 2 else 1 else: batch_size = len(prompt_ids) if prompt_ids and isinstance(prompt_ids[0], list) else 1 has_neg_prompt = negative_prompt_ids is not None do_true_cfg = true_cfg_scale > 1 and has_neg_prompt self.check_cfg_parallel_validity(true_cfg_scale, has_neg_prompt) prompt_embeds, prompt_embeds_mask = self.encode_prompt( prompt_ids=prompt_ids, attention_mask=prompt_mask, num_images_per_prompt=num_images_per_prompt, max_sequence_length=max_sequence_length, ) if do_true_cfg: negative_prompt_embeds, negative_prompt_embeds_mask = self.encode_prompt( prompt_ids=negative_prompt_ids, attention_mask=negative_prompt_mask, num_images_per_prompt=num_images_per_prompt, max_sequence_length=max_sequence_length, ) else: negative_prompt_embeds = None negative_prompt_embeds_mask = None num_channels_latents = self.transformer.in_channels // 4 latents = self.prepare_latents( batch_size * num_images_per_prompt, num_channels_latents, height, width, torch.float32, self.device, generator, None, ) img_shapes = build_img_shapes(height, width, batch_size, self.vae_scale_factor) timesteps, _ = self.prepare_timesteps(num_inference_steps, sigmas, latents.shape[1]) self._num_timesteps = len(timesteps) if self.transformer.guidance_embeds: guidance = torch.full([1], guidance_scale, dtype=torch.float32) guidance = guidance.expand(latents.shape[0]) else: guidance = None txt_seq_lens = txt_seq_lens_from_embeds(prompt_embeds) negative_txt_seq_lens = txt_seq_lens_from_embeds(negative_prompt_embeds) req_scheduler = copy.deepcopy(self.scheduler) req_scheduler.set_begin_index(0) # Resolve SDE / log-prob knobs from sampling extra_args so that the # step-execution path mirrors ``forward()``'s rollout behaviour. extra = sampling.extra_args or {} noise_level = coalesce_not_none(extra.get("noise_level", None), 0.7) sde_window_size = coalesce_not_none(extra.get("sde_window_size", None), None) sde_window_range = coalesce_not_none(extra.get("sde_window_range", None), (0, 5)) sde_type = coalesce_not_none(extra.get("sde_type", None), "sde") logprobs = coalesce_not_none(extra.get("logprobs", None), True) if sde_window_size is not None: start = torch.randint( sde_window_range[0], sde_window_range[1] - sde_window_size + 1, (1,), generator=generator, device=self.device, ).item() sde_window = (start, start + sde_window_size) else: sde_window = (0, len(timesteps) - 1) state.prompt_embeds = prompt_embeds state.prompt_embeds_mask = prompt_embeds_mask state.negative_prompt_embeds = negative_prompt_embeds state.negative_prompt_embeds_mask = negative_prompt_embeds_mask state.latents = latents state.timesteps = timesteps state.step_index = 0 state.scheduler = req_scheduler state.do_true_cfg = do_true_cfg state.guidance = guidance state.img_shapes = img_shapes state.txt_seq_lens = txt_seq_lens state.negative_txt_seq_lens = negative_txt_seq_lens state.sampling.cfg_normalize = True # Persist the resolved generator so ``step_scheduler`` (executed # one step at a time by the step-execution engine) keeps drawing # from the same RNG stream as ``forward()``. state.sampling.generator = generator # Rollout / SDE state consumed by ``step_scheduler`` and packaged # into trajectory/metadata fields by ``post_decode``. state.sde_window = sde_window state.noise_level = noise_level state.sde_type = sde_type state.logprobs = logprobs state.all_latents = [] state.all_log_probs = [] state.all_timesteps = [] return state def diffuse( self, prompt_embeds, prompt_embeds_mask, negative_prompt_embeds, negative_prompt_embeds_mask, latents, img_shapes, txt_seq_lens, negative_txt_seq_lens, timesteps, do_true_cfg, guidance, true_cfg_scale, noise_level, sde_window, sde_type, generator, logprobs, ): """Run the full SDE diffusion loop and collect per-step rollout data. Iterates over all timesteps, optionally applying True-CFG guidance, and collects latents and log-probabilities within the SDE window. Args: prompt_embeds (torch.Tensor): Positive prompt embeddings. prompt_embeds_mask (torch.Tensor): Attention mask for *prompt_embeds*. negative_prompt_embeds (torch.Tensor): Negative prompt embeddings for CFG. negative_prompt_embeds_mask (torch.Tensor): Attention mask for *negative_prompt_embeds*. latents (torch.Tensor): Initial noisy latents. img_shapes (list): Per-sample image shapes used by the transformer. txt_seq_lens (list[int]): Sequence lengths for positive prompt embeddings. negative_txt_seq_lens (list[int]): Sequence lengths for negative prompt embeddings. timesteps (torch.Tensor): Scheduler timestep sequence. do_true_cfg (bool): Whether to apply True-CFG guidance. guidance (torch.Tensor | None): Guidance scale tensor, or ``None``. true_cfg_scale (float): Classifier-free guidance scale. noise_level (float): SDE noise injection magnitude within the window. sde_window (tuple[int, int] | list[tuple[int, int]]): Shared or per-row ``(start, end)`` window(s). sde_type (str): SDE variant; one of ``"sde"`` or ``"cps"``. generator (torch.Generator | list[torch.Generator] | None): Optional RNG; a list must have one entry per batch row. logprobs (bool): Whether to compute and return per-step log-probabilities. Returns: tuple: A 4-tuple of ``(latents, all_latents, all_log_probs, all_timesteps)`` where *all_latents* has shape ``(B, W+1, ...)`` (W = SDE-window length), *all_log_probs* has shape ``(B, W)`` or ``None`` when *logprobs* is ``False``, and *all_timesteps* has shape ``(B, W)``. """ batch_size = latents.shape[0] windows = [sde_window] * batch_size if isinstance(sde_window, tuple) else list(sde_window) if len(windows) != batch_size: raise ValueError(f"Expected {batch_size} SDE windows, got {len(windows)}.") if len({end - start for start, end in windows}) != 1: raise ValueError("Packed SDE windows must share the same size.") all_latents: list[list[torch.Tensor]] = [[] for _ in range(batch_size)] all_log_probs: list[list[Any]] = [[] for _ in range(batch_size)] all_timesteps: list[list[Any]] = [[] for _ in range(batch_size)] self.scheduler.set_begin_index(0) for i, timestep_value in enumerate(timesteps): if self.interrupt: continue for batch_idx, (start, end) in enumerate(windows): if i == start: all_latents[batch_idx].append(latents[batch_idx].detach().float().clone()) levels = [float(noise_level) if start <= i < end else 0.0 for start, end in windows] cur_noise_level: float | torch.Tensor = ( levels[0] if all(level == levels[0] for level in levels) else torch.tensor(levels, device=latents.device, dtype=torch.float32).view( batch_size, *([1] * (latents.ndim - 1)) ) ) self._current_timestep = timestep_value # Broadcast timestep to match batch size timestep = timestep_value.expand(latents.shape[0]).to(device=latents.device, dtype=latents.dtype) # Cast to model dtype for transformer forward (scheduler returns float32). x = latents.to(self.transformer.img_in.weight.dtype) self.transformer.do_true_cfg = do_true_cfg # Forward pass for positive prompt (or unconditional if no CFG) noise_pred = self.transformer( hidden_states=x, timestep=timestep / 1000, guidance=guidance, encoder_hidden_states_mask=prompt_embeds_mask, encoder_hidden_states=prompt_embeds, img_shapes=img_shapes, txt_seq_lens=txt_seq_lens, attention_kwargs=self.attention_kwargs, return_dict=False, )[0] # Forward pass for negative prompt (CFG) if do_true_cfg: neg_noise_pred = self.transformer( hidden_states=x, timestep=timestep / 1000, guidance=guidance, encoder_hidden_states_mask=negative_prompt_embeds_mask, encoder_hidden_states=negative_prompt_embeds, img_shapes=img_shapes, txt_seq_lens=negative_txt_seq_lens, attention_kwargs=self.attention_kwargs, return_dict=False, )[0] noise_pred = apply_true_cfg(noise_pred, neg_noise_pred, true_cfg_scale) # compute the previous noisy sample x_t -> x_t-1 latents, log_prob, _, _ = self.scheduler.step( noise_pred.to(torch.float32), timestep_value, latents.to(torch.float32), generator=generator, noise_level=cur_noise_level, sde_type=sde_type, return_logprobs=logprobs, return_dict=False, ) # Save fp32 trajectory BEFORE casting to model dtype, so the # trainer recomputes log-probs on full-precision latents. for batch_idx, (start, end) in enumerate(windows): if start <= i < end: all_latents[batch_idx].append(latents[batch_idx].detach().to(torch.float32).clone()) all_log_probs[batch_idx].append(None if log_prob is None else log_prob[batch_idx]) all_timesteps[batch_idx].append(timestep_value) all_latents_t = torch.stack([torch.stack(traj, dim=0) for traj in all_latents], dim=0) if all_log_probs and all_log_probs[0] and all_log_probs[0][0] is not None: all_log_probs_t = torch.stack([torch.stack(traj, dim=0) for traj in all_log_probs], dim=0) else: all_log_probs_t = None all_timesteps_t = torch.stack([torch.stack(traj, dim=0) for traj in all_timesteps], dim=0) return latents, all_latents_t, all_log_probs_t, all_timesteps_t def step_scheduler( self, state: StepRequestState, noise_pred: torch.Tensor, **kwargs: Any, ) -> None: """One scheduler step that mirrors the per-iter body of :meth:`diffuse`. The default ``QwenImagePipeline.step_scheduler`` calls the standard scheduler.step without SDE noise / log-prob bookkeeping, which means ``step_execution=True`` would silently drop ``all_latents`` / ``all_log_probs`` / ``all_timesteps`` (and the ``prompt_embeds_mask`` consumer downstream would then receive a ``None`` value that turns into a non-tensor ``LinkedList`` inside the training ``TensorDict``). Override here to keep the step-mode and request-mode trajectories equivalent. """ del kwargs if self.interrupt: return i = state.step_index timestep_value = state.timesteps[i] sde_window = state.sde_window if i < sde_window[0]: cur_noise_level = 0.0 elif i == sde_window[0]: cur_noise_level = state.noise_level state.all_latents.append(state.latents.to(torch.float32)) elif i > sde_window[0] and i < sde_window[1]: cur_noise_level = state.noise_level else: cur_noise_level = 0.0 new_latents, log_prob, _, _ = state.scheduler.step( noise_pred.to(torch.float32), timestep_value, state.latents.to(torch.float32), generator=state.sampling.generator, noise_level=cur_noise_level, sde_type=state.sde_type, return_logprobs=state.logprobs, return_dict=False, ) # Save fp32 trajectory so the trainer later recomputes log-probs on # full-precision latents. if i >= sde_window[0] and i < sde_window[1]: state.all_latents.append(new_latents.to(torch.float32)) state.all_log_probs.append(log_prob) state.all_timesteps.append(timestep_value) # Keep live state in fp32 for the whole trajectory. ``denoise_step`` # already casts latents to the transformer dtype before the forward # pass, so storing model_dtype here is unnecessary. More importantly, # under continuous batching the engine gathers ``state.latents`` across # all in-flight requests: a freshly-added request still holds fp32 # latents from ``prepare_encode`` while stepped requests would hold # model-dtype latents, producing a "Mixed dtypes in latents batch" # error. Keeping fp32 throughout makes the batch dtype consistent. state.latents = new_latents.to(torch.float32) state.step_index += 1 def denoise_step(self, input_batch, **kwargs): del kwargs if self.interrupt: return None t = input_batch.timesteps self._current_timestep = t x = input_batch.latents.to(self.transformer.img_in.weight.dtype) positive_kwargs, negative_kwargs, output_slice = self._build_denoise_kwargs( latents=x, timestep=t, guidance=input_batch.guidance, prompt_embeds=input_batch.prompt_embeds, prompt_embeds_mask=input_batch.prompt_embeds_mask, img_shapes=input_batch.img_shapes, txt_seq_lens=input_batch.txt_seq_lens, do_true_cfg=input_batch.do_true_cfg, negative_prompt_embeds=input_batch.negative_prompt_embeds, negative_prompt_embeds_mask=input_batch.negative_prompt_embeds_mask, negative_txt_seq_lens=input_batch.negative_txt_seq_lens, extra_transformer_kwargs={"attention_kwargs": self.attention_kwargs, "return_dict": False}, ) noise_pred = self.predict_noise_maybe_with_cfg( input_batch.do_true_cfg, input_batch.true_cfg_scale, positive_kwargs, negative_kwargs, input_batch.cfg_normalize, output_slice, ) return noise_pred.float() # step_scheduler expects fp32 noise_pred def post_decode( self, state: StepRequestState, **kwargs: Any, ) -> DiffusionOutput: """Decode final latents, package rollout trajectory, and move to CPU. In ``step_execution`` mode the worker ships the returned :class:`DiffusionOutput` across an inter-process MessageQueue to the ``vLLMOmniHttpServer`` actor. We must (a) move tensors to CPU so the receiving process does not initialise a stray CUDA context on GPU 0, and (b) populate native ``trajectory_*`` / metadata fields that :meth:`forward` produces, so downstream consumers (``vllm_omni_async_server.generate`` -> ``embeds_padding_2_no_padding``) receive real tensors rather than ``None`` (which becomes a non-tensor ``LinkedList`` in the ``TensorDict`` and breaks ``mask.shape[0]``). """ output = super().post_decode(state, **kwargs) if not isinstance(output, DiffusionOutput): return output all_latents = state.all_latents all_log_probs = state.all_log_probs all_timesteps = state.all_timesteps stacked_latents = torch.stack(all_latents, dim=1) if all_latents else None stacked_log_probs = ( torch.stack(all_log_probs, dim=1) if all_log_probs and all_log_probs[0] is not None else None ) stacked_timesteps = ( torch.stack(all_timesteps).unsqueeze(0).expand(state.latents.shape[0], -1) if all_timesteps else None ) return with_rollout_data( output, trajectory_latents=stacked_latents, trajectory_log_probs=stacked_log_probs, trajectory_timesteps=stacked_timesteps, prompt_embeddings={ "prompt_embeds": state.prompt_embeds, "prompt_embeds_mask": state.prompt_embeds_mask, "negative_prompt_embeds": state.negative_prompt_embeds, "negative_prompt_embeds_mask": state.negative_prompt_embeds_mask, }, to_cpu=True, ) def forward( self, req: OmniDiffusionRequest | DiffusionRequestBatch, prompt_token_ids: torch.Tensor | list[int] | list[list[int]] | None = None, prompt_mask: torch.Tensor | list[int] | list[bool] | list[list[int]] | list[list[bool]] | None = None, negative_prompt_ids: torch.Tensor | list[int] | list[list[int]] | None = None, negative_prompt_mask: torch.Tensor | list[int] | list[bool] | list[list[int]] | list[list[bool]] | None = None, true_cfg_scale: float = 4.0, height: int | None = None, width: int | None = None, num_inference_steps: int = 50, sigmas: list[float] | None = None, guidance_scale: float = 1.0, num_images_per_prompt: int = 1, generator: torch.Generator | list[torch.Generator] | None = None, latents: torch.Tensor | None = None, prompt_embeds: torch.Tensor | None = None, prompt_embeds_mask: torch.Tensor | None = None, negative_prompt_embeds: torch.Tensor | None = None, negative_prompt_embeds_mask: torch.Tensor | None = None, output_type: str | None = "pil", attention_kwargs: dict[str, Any] | None = None, callback_on_step_end_tensor_inputs: tuple[str, ...] = ("latents",), max_sequence_length: int = 512, noise_level: float = 0.7, sde_window_size: int | None = None, sde_window_range: tuple[int, int] = (0, 5), sde_type: Literal["sde", "cps"] = "sde", logprobs: bool = True, ) -> DiffusionOutput | list[DiffusionOutput]: """End-to-end image generation with rollout data collection. Encodes the prompt, prepares latents, runs the SDE diffusion loop via :meth:`diffuse`, and decodes the final latents through the VAE. Sampling parameters in *req* take precedence over the keyword arguments. Args: req (OmniDiffusionRequest | DiffusionRequestBatch): One rollout request or a request batch containing prompts and sampling parameters. prompt_token_ids (torch.Tensor | list[int] | list[list[int]], *optional*): Token IDs for the positive prompt. prompt_mask (torch.Tensor | list, *optional*): Attention mask for *prompt_token_ids*. negative_prompt_ids (torch.Tensor | list[int] | list[list[int]], *optional*): Token IDs for the negative prompt used in True-CFG. negative_prompt_mask (torch.Tensor | list, *optional*): Attention mask for *negative_prompt_ids*. true_cfg_scale (float): Classifier-free guidance scale; CFG is disabled when ``<= 1``. height (int, *optional*): Output image height in pixels. width (int, *optional*): Output image width in pixels. num_inference_steps (int): Number of denoising steps. sigmas (list[float], *optional*): Custom sigmas for the scheduler. guidance_scale (float): Distilled guidance scale embedded in the transformer (``guidance_embeds`` mode). num_images_per_prompt (int): Number of images to generate per prompt. generator (torch.Generator | list[torch.Generator], *optional*): Random generator(s) for reproducibility. latents (torch.Tensor, *optional*): Pre-generated initial latents; sampled from a Gaussian when ``None``. prompt_embeds (torch.Tensor, *optional*): Pre-computed positive prompt embeddings; bypasses the text encoder. prompt_embeds_mask (torch.Tensor, *optional*): Attention mask for pre-computed *prompt_embeds*. negative_prompt_embeds (torch.Tensor, *optional*): Pre-computed negative prompt embeddings. negative_prompt_embeds_mask (torch.Tensor, *optional*): Attention mask for *negative_prompt_embeds*. output_type (str, *optional*): Format of the returned image; ``"latent"`` returns raw latents, otherwise the VAE-decoded image. attention_kwargs (dict, *optional*): Extra keyword arguments forwarded to the attention layers. callback_on_step_end_tensor_inputs (tuple[str, ...]): Names of tensors to expose in the step-end callback. max_sequence_length (int): Maximum prompt embedding sequence length. noise_level (float): SDE noise injection magnitude within the window. sde_window_size (int, *optional*): Number of SDE steps; when ``None`` the full timestep range is used. sde_window_range (tuple[int, int]): ``(start, end)`` range from which the SDE window start position is randomly sampled. sde_type (str): SDE variant; ``"sde"`` or ``"cps"``. logprobs (bool): Whether to compute per-step log-probabilities. Returns: DiffusionOutput | list[DiffusionOutput]: Contains the decoded *output* image plus native ``trajectory_*`` fields and prompt-embedding metadata used for FlowGRPO training. """ request_batch = req if isinstance(req, DiffusionRequestBatch) else DiffusionRequestBatch(requests=[req]) return_batch = isinstance(req, DiffusionRequestBatch) prompts = request_batch.prompts prompt_embed_cache = getattr(self, "_prompt_embed_cache", None) use_prompt_embed_cache = bool(prompt_embed_cache is not None and prompt_embed_cache.enabled) # The prompt cache wraps encode_prompt, so keep token inputs as lists until that boundary. prompt_token_ids, prompt_token_lengths = _collate_prompt_rows( prompts, ("prompt_token_ids", "prompt_ids"), prompt_token_ids, device=self.device, field_name="prompt_token_ids", preserve_lists=use_prompt_embed_cache, ) prompt_mask = _collate_prompt_mask( prompts, ("prompt_mask",), prompt_mask, device=self.device, field_name="prompt_mask", token_lengths=prompt_token_lengths, target_seq_len=max(prompt_token_lengths) if prompt_token_lengths else None, preserve_lists=use_prompt_embed_cache, ) negative_prompt_ids, negative_prompt_lengths = _collate_prompt_rows( prompts, ("negative_prompt_ids",), negative_prompt_ids, device=self.device, field_name="negative_prompt_ids", preserve_lists=use_prompt_embed_cache, ) negative_prompt_mask = _collate_prompt_mask( prompts, ("negative_prompt_mask",), negative_prompt_mask, device=self.device, field_name="negative_prompt_mask", token_lengths=negative_prompt_lengths, target_seq_len=max(negative_prompt_lengths) if negative_prompt_lengths else None, preserve_lists=use_prompt_embed_cache, ) sampling_params = request_batch.sampling_params_list[0] height = sampling_params.height or self.default_sample_size * self.vae_scale_factor width = sampling_params.width or self.default_sample_size * self.vae_scale_factor num_inference_steps = sampling_params.num_inference_steps or num_inference_steps sigmas = sampling_params.sigmas or sigmas max_sequence_length = sampling_params.max_sequence_length or max_sequence_length output_type = sampling_params.output_type or output_type noise_level = coalesce_not_none(sampling_params.extra_args.get("noise_level", None), noise_level) sde_window_size = coalesce_not_none(sampling_params.extra_args.get("sde_window_size", None), sde_window_size) sde_window_range = coalesce_not_none(sampling_params.extra_args.get("sde_window_range", None), sde_window_range) sde_type = coalesce_not_none(sampling_params.extra_args.get("sde_type", None), sde_type) logprobs = coalesce_not_none(sampling_params.extra_args.get("logprobs", None), logprobs) for request in request_batch.requests: request_sampling_params = request.sampling_params if request_sampling_params.generator is None and request_sampling_params.seed is not None: request_sampling_params.generator = torch.Generator(device=self.device).manual_seed( request_sampling_params.seed ) true_cfg_scale = coalesce_not_none(sampling_params.true_cfg_scale, true_cfg_scale) if getattr(sampling_params, "guidance_scale_provided", False): guidance_scale = sampling_params.guidance_scale req_num_outputs = getattr(sampling_params, "num_outputs_per_prompt", None) if req_num_outputs and req_num_outputs > 0: num_images_per_prompt = req_num_outputs generator = request_batch.collate_request_generators(num_images_per_prompt, generator) latents = request_batch.collate_request_tensors("latents", latents) self._guidance_scale = guidance_scale self._attention_kwargs = attention_kwargs self._current_timestep = None self._interrupt = False if prompt_token_ids is not None: if isinstance(prompt_token_ids, torch.Tensor): batch_size = prompt_token_ids.shape[0] if prompt_token_ids.ndim == 2 else 1 else: batch_size = len(prompt_token_ids) elif prompt_embeds is not None: batch_size = prompt_embeds.shape[0] else: # Both prompt_token_ids and prompt_embeds are None (e.g. during warmup/dummy run). # Return a minimal dummy output to avoid crashing. outputs = [DiffusionOutput(output=None) for _ in range(request_batch.num_reqs)] return outputs if return_batch else outputs[0] has_neg_prompt = negative_prompt_ids is not None or ( negative_prompt_embeds is not None and negative_prompt_embeds_mask is not None ) do_true_cfg = true_cfg_scale > 1 and has_neg_prompt prompt_embeds, prompt_embeds_mask = self.encode_prompt( prompt_ids=prompt_token_ids, attention_mask=prompt_mask, prompt_embeds=prompt_embeds, prompt_embeds_mask=prompt_embeds_mask, num_images_per_prompt=num_images_per_prompt, max_sequence_length=max_sequence_length, ) if do_true_cfg: negative_prompt_embeds, negative_prompt_embeds_mask = self.encode_prompt( prompt_ids=negative_prompt_ids, attention_mask=negative_prompt_mask, prompt_embeds=negative_prompt_embeds, prompt_embeds_mask=negative_prompt_embeds_mask, num_images_per_prompt=num_images_per_prompt, max_sequence_length=max_sequence_length, ) num_channels_latents = self.transformer.in_channels // 4 latents = self.prepare_latents( batch_size * num_images_per_prompt, num_channels_latents, height, width, prompt_embeds.dtype, self.device, generator, latents, ) img_shapes = build_img_shapes(height, width, batch_size, self.vae_scale_factor) timesteps, num_inference_steps = self.prepare_timesteps(num_inference_steps, sigmas, latents.shape[1]) self._num_timesteps = len(timesteps) if self.transformer.guidance_embeds: guidance = torch.full([1], guidance_scale, dtype=torch.float32) guidance = guidance.expand(latents.shape[0]) else: guidance = None if self.attention_kwargs is None: self._attention_kwargs = {} txt_seq_lens = prompt_embeds_mask.sum(dim=1).tolist() if prompt_embeds_mask is not None else None negative_txt_seq_lens = ( negative_prompt_embeds_mask.sum(dim=1).tolist() if negative_prompt_embeds_mask is not None else None ) sde_window = _sample_per_sample_sde_windows( sde_window_size=sde_window_size, sde_window_range=sde_window_range if sde_window_range is not None else (0, 5), num_timesteps=len(timesteps), batch_size=latents.shape[0], generator=generator, device=self.device, ) latents, all_latents, all_log_probs, all_timesteps = self.diffuse( prompt_embeds, prompt_embeds_mask, negative_prompt_embeds, negative_prompt_embeds_mask, latents, img_shapes, txt_seq_lens, negative_txt_seq_lens, timesteps, do_true_cfg, guidance, true_cfg_scale, noise_level, sde_window, sde_type, generator, logprobs, ) self._current_timestep = None if output_type == "latent": image = latents else: latents = self._unpack_latents(latents, height, width, self.vae_scale_factor) latents = latents.to(self.vae.dtype) latents_mean = ( torch.tensor(self.vae.config.latents_mean) .view(1, self.vae.config.z_dim, 1, 1, 1) .to(latents.device, latents.dtype) ) latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to( latents.device, latents.dtype ) latents = latents / latents_std + latents_mean image = self.vae.decode(latents, return_dict=False)[0][:, :, 0] result = rollout_output( media=image, trajectory_latents=all_latents, trajectory_log_probs=all_log_probs, trajectory_timesteps=all_timesteps, prompt_embeddings={ "prompt_embeds": prompt_embeds, "prompt_embeds_mask": prompt_embeds_mask, "negative_prompt_embeds": negative_prompt_embeds, "negative_prompt_embeds_mask": negative_prompt_embeds_mask, }, to_cpu=True, ) outputs = _split_diffusion_output_by_request( result, request_batch, num_outputs_per_prompt=num_images_per_prompt, ) return outputs if return_batch else outputs[0]