689 lines
28 KiB
Python
689 lines
28 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
"""Test doubles for custom-pipeline E2E and RLHF integration.
|
|
|
|
Three layers:
|
|
|
|
* ``FlowMatchSDEDiscreteSchedulerForTest`` — SDE scheduler with per-step logprobs.
|
|
* ``vLLMOmniColocateWorkerExtensionForTest`` — verl-omni colocate worker extension mirror.
|
|
* ``QwenImagePipelineWithLogProbForTest`` — Qwen-Image pipeline returning trajectory_* fields.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import ctypes
|
|
import logging
|
|
import math
|
|
import os
|
|
import platform
|
|
import signal
|
|
from dataclasses import dataclass
|
|
from typing import Any, Literal
|
|
|
|
import torch
|
|
from diffusers import FlowMatchEulerDiscreteScheduler
|
|
from diffusers.utils import BaseOutput
|
|
from diffusers.utils.torch_utils import randn_tensor
|
|
from msgspec import field
|
|
|
|
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.worker.diffusion_worker import CustomPipelineWorkerExtension
|
|
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
|
|
from vllm_omni.lora.request import LoRARequest as OmniLoRARequest
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Scheduler layer
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@dataclass
|
|
class FlowMatchSDEDiscreteSchedulerOutputForTest(BaseOutput):
|
|
"""
|
|
Output class for the scheduler's `step` function output.
|
|
|
|
Args:
|
|
prev_sample (`torch.FloatTensor` of shape `(batch_size, sequence_length, num_channels)` for images):
|
|
Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the
|
|
denoising loop.
|
|
log_prob (`torch.FloatTensor` of shape `(batch_size,)`, *optional*):
|
|
The log probability of the previous sample.
|
|
prev_sample_mean (`torch.FloatTensor` of shape `(batch_size, sequence_length, num_channels)` for images):
|
|
The mean of the computed sample of previous timestep.
|
|
std_dev_t (`torch.FloatTensor` of shape `(batch_size, 1, 1)`):
|
|
The standard deviation used to compute `prev_sample`.
|
|
"""
|
|
|
|
prev_sample: torch.FloatTensor
|
|
log_prob: torch.FloatTensor | None
|
|
prev_sample_mean: torch.FloatTensor
|
|
std_dev_t: torch.FloatTensor
|
|
|
|
|
|
class FlowMatchSDEDiscreteSchedulerForTest(FlowMatchEulerDiscreteScheduler):
|
|
"""SDE version of the FlowMatchEulerDiscreteScheduler.
|
|
The implementation is based on FlowGRPO paper (https://arxiv.org/abs/2505.05470)
|
|
and diffusers v0.37 branch.
|
|
"""
|
|
|
|
def step(
|
|
self,
|
|
model_output: torch.FloatTensor,
|
|
timestep: 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,
|
|
generator: torch.Generator | None = None,
|
|
per_token_timesteps: torch.Tensor | None = None,
|
|
return_dict: bool = True,
|
|
noise_level: float = 0.7,
|
|
prev_sample: torch.FloatTensor | None = None,
|
|
sde_type: Literal["sde", "cps"] = "sde",
|
|
return_logprobs: bool = True,
|
|
) -> FlowMatchSDEDiscreteSchedulerOutputForTest | 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).
|
|
|
|
Modified from https://github.com/yifan123/flow_grpo/blob/main/flow_grpo/diffusers_patch/sd3_sde_with_logprob.py
|
|
|
|
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.
|
|
per_token_timesteps (`torch.Tensor`, *optional*):
|
|
The timesteps for each token in the sample.
|
|
return_dict (`bool`):
|
|
Whether or not to return a
|
|
[`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchSDEDiscreteSchedulerOutput`] or tuple.
|
|
noise_level (`float`, *optional*, defaults to 0.7):
|
|
The noise level used in the SDE.
|
|
prev_sample (`torch.FloatTensor`, *optional*):
|
|
The sample from the previous timestep. If not provided, it will be sampled inside the function.
|
|
sde_type (`str`, *optional*, defaults to "sde"):
|
|
The type of SDE to use. Choose between "sde" and "cps".
|
|
return_logprobs (`bool`, *optional*, defaults to True):
|
|
Whether to return log probabilities of the previous sample.
|
|
"""
|
|
|
|
if isinstance(timestep, (int, torch.IntTensor, torch.LongTensor)):
|
|
raise ValueError(
|
|
(
|
|
"Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to"
|
|
" `FlowMatchEulerDiscreteScheduler.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
|
|
sample = sample.to(torch.float32)
|
|
if prev_sample is not None:
|
|
prev_sample = prev_sample.to(torch.float32)
|
|
|
|
prev_sample, log_prob, prev_sample_mean, std_dev_t = self.sample_previous_step(
|
|
sample=sample,
|
|
model_output=model_output,
|
|
generator=generator,
|
|
per_token_timesteps=per_token_timesteps,
|
|
noise_level=noise_level,
|
|
prev_sample=prev_sample,
|
|
sde_type=sde_type,
|
|
return_logprobs=return_logprobs,
|
|
)
|
|
|
|
# upon completion increase step index by one
|
|
self._step_index += 1
|
|
if per_token_timesteps is None:
|
|
# Cast sample back to model compatible dtype
|
|
prev_sample = prev_sample.to(model_output.dtype)
|
|
|
|
if not return_dict:
|
|
return (prev_sample, log_prob, prev_sample_mean, std_dev_t)
|
|
|
|
return FlowMatchSDEDiscreteSchedulerOutputForTest(
|
|
prev_sample=prev_sample, log_prob=log_prob, prev_sample_mean=prev_sample_mean, std_dev_t=std_dev_t
|
|
)
|
|
|
|
def sample_previous_step(
|
|
self,
|
|
sample: torch.Tensor,
|
|
model_output: torch.Tensor,
|
|
timestep: torch.FloatTensor | None = None,
|
|
generator: torch.Generator | None = None,
|
|
per_token_timesteps: torch.Tensor | None = None,
|
|
noise_level: float = 0.7,
|
|
prev_sample: torch.Tensor | None = None,
|
|
sde_type: Literal["cps", "sde"] = "sde",
|
|
return_logprobs: bool = True,
|
|
):
|
|
assert sde_type in ["sde", "cps"]
|
|
assert sample.dtype == torch.float32
|
|
if prev_sample is not None:
|
|
assert prev_sample.dtype == torch.float32
|
|
|
|
if per_token_timesteps is not None:
|
|
raise NotImplementedError("per_token_timesteps is not supported yet for FlowMatchSDEDiscreteScheduler.")
|
|
else:
|
|
if timestep is None:
|
|
sigma_idx = self.step_index
|
|
sigma = self.sigmas[sigma_idx]
|
|
sigma_prev = self.sigmas[sigma_idx + 1]
|
|
else:
|
|
sigma_idx = torch.tensor([self.index_for_timestep(t) for t in timestep])
|
|
sigma = self.sigmas[sigma_idx].view(-1, *([1] * (len(sample.shape) - 1)))
|
|
sigma_prev = self.sigmas[sigma_idx + 1].view(-1, *([1] * (len(sample.shape) - 1)))
|
|
|
|
sigma_max = self.sigmas[1]
|
|
dt = sigma_prev - sigma
|
|
|
|
if sde_type == "sde":
|
|
std_dev_t = torch.sqrt(sigma / (1 - torch.where(sigma == 1, sigma_max, sigma))) * noise_level
|
|
|
|
prev_sample_mean = (
|
|
sample * (1 + std_dev_t**2 / (2 * sigma) * dt)
|
|
+ model_output * (1 + std_dev_t**2 * (1 - sigma) / (2 * sigma)) * dt
|
|
)
|
|
|
|
if prev_sample is None:
|
|
variance_noise = randn_tensor(
|
|
model_output.shape,
|
|
generator=generator,
|
|
device=model_output.device,
|
|
dtype=model_output.dtype,
|
|
)
|
|
prev_sample = prev_sample_mean + std_dev_t * torch.sqrt(-1 * dt) * variance_noise
|
|
|
|
if return_logprobs:
|
|
log_prob = (
|
|
-((prev_sample.detach() - prev_sample_mean) ** 2) / (2 * ((std_dev_t * torch.sqrt(-1 * dt)) ** 2))
|
|
- torch.log(std_dev_t * torch.sqrt(-1 * dt))
|
|
- torch.log(torch.sqrt(2 * torch.as_tensor(math.pi)))
|
|
)
|
|
else:
|
|
log_prob = None
|
|
|
|
elif sde_type == "cps":
|
|
std_dev_t = sigma_prev * math.sin(noise_level * math.pi / 2)
|
|
pred_original_sample = sample - sigma * model_output
|
|
noise_estimate = sample + model_output * (1 - sigma)
|
|
prev_sample_mean = pred_original_sample * (1 - sigma_prev) + noise_estimate * torch.sqrt(
|
|
sigma_prev**2 - std_dev_t**2
|
|
)
|
|
|
|
if prev_sample is None:
|
|
variance_noise = randn_tensor(
|
|
model_output.shape,
|
|
generator=generator,
|
|
device=model_output.device,
|
|
dtype=model_output.dtype,
|
|
)
|
|
prev_sample = prev_sample_mean + std_dev_t * variance_noise
|
|
|
|
if return_logprobs:
|
|
log_prob = -((prev_sample.detach() - prev_sample_mean) ** 2)
|
|
else:
|
|
log_prob = None
|
|
|
|
# mean along all but batch dimension
|
|
log_prob = log_prob.mean(dim=tuple(range(1, log_prob.ndim))) if log_prob is not None else None
|
|
return prev_sample, log_prob, prev_sample_mean, std_dev_t
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Worker extension layer
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class OmniTensorLoRARequestForTest(OmniLoRARequest):
|
|
peft_config: dict = field(default=None)
|
|
lora_tensors: dict = field(default=None)
|
|
|
|
|
|
def set_death_signal() -> None:
|
|
"""Verbatim port of ``verl.workers.rollout.vllm_rollout.utils.set_death_signal``."""
|
|
if platform.system() != "Linux":
|
|
return
|
|
try:
|
|
libc = ctypes.CDLL("libc.so.6")
|
|
libc.prctl(1, signal.SIGKILL)
|
|
if os.getppid() == 1:
|
|
os.kill(os.getpid(), signal.SIGKILL)
|
|
except Exception: # noqa: BLE001
|
|
# libc.so.6 may not be present (musl, alpine) — best-effort only.
|
|
pass
|
|
|
|
|
|
class vLLMOmniColocateWorkerExtensionForTest(CustomPipelineWorkerExtension):
|
|
"""Mirror of ``vLLMOmniColocateWorkerExtension`` (verl-omni).
|
|
|
|
The production ``__new__`` runs ``set_death_signal`` + ``VLLMOmniHijack.hijack``
|
|
on every vLLM worker process. Replicating both ensures this test
|
|
reproduces the same monkey-patched environment.
|
|
"""
|
|
|
|
def __new__(cls, **kwargs):
|
|
set_death_signal()
|
|
return super().__new__(cls)
|
|
|
|
@staticmethod
|
|
def test_extension_name() -> str:
|
|
"""Return a stable identifier for assertions in unit tests."""
|
|
return "vllm-omni-colocate-worker-extension-for-test"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Pipeline layer
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _maybe_to_cpu(v):
|
|
if isinstance(v, torch.Tensor):
|
|
return v.detach().cpu()
|
|
return v
|
|
|
|
|
|
class QwenImagePipelineWithLogProbForTest(QwenImagePipeline):
|
|
supports_request_batch = False
|
|
|
|
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
|
|
# Check if model is a local path
|
|
local_files_only = os.path.exists(model)
|
|
|
|
self.scheduler = FlowMatchSDEDiscreteSchedulerForTest.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,
|
|
attention_mask: torch.Tensor | 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,
|
|
):
|
|
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
|
|
)
|
|
batch_size = prompt_ids.shape[0] if prompt_embeds is None else prompt_embeds.shape[0]
|
|
|
|
if prompt_embeds is None:
|
|
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]
|
|
|
|
_, seq_len, _ = prompt_embeds.shape
|
|
prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
|
|
prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
|
|
prompt_embeds_mask = prompt_embeds_mask.repeat(1, num_images_per_prompt, 1)
|
|
prompt_embeds_mask = prompt_embeds_mask.view(batch_size * num_images_per_prompt, seq_len)
|
|
|
|
return prompt_embeds, prompt_embeds_mask
|
|
|
|
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,
|
|
):
|
|
all_latents = []
|
|
all_log_probs = []
|
|
all_timesteps = []
|
|
self.scheduler.set_begin_index(0)
|
|
for i, t in enumerate(timesteps):
|
|
if self.interrupt:
|
|
continue
|
|
|
|
if i < sde_window[0]:
|
|
cur_noise_level = 0.0
|
|
elif i == sde_window[0]:
|
|
cur_noise_level = noise_level
|
|
all_latents.append(latents)
|
|
elif i > sde_window[0] and i < sde_window[1]:
|
|
cur_noise_level = noise_level
|
|
else:
|
|
cur_noise_level = 0.0
|
|
|
|
self._current_timestep = t
|
|
|
|
# Broadcast timestep to match batch size
|
|
timestep = t.expand(latents.shape[0]).to(device=latents.device, dtype=latents.dtype)
|
|
|
|
# Forward pass for positive prompt (or unconditional if no CFG)
|
|
self.transformer.do_true_cfg = do_true_cfg
|
|
noise_pred = self.transformer(
|
|
hidden_states=latents,
|
|
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=latents,
|
|
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]
|
|
comb_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred)
|
|
cond_norm = torch.norm(noise_pred, dim=-1, keepdim=True)
|
|
noise_norm = torch.norm(comb_pred, dim=-1, keepdim=True)
|
|
noise_pred = comb_pred * (cond_norm / noise_norm)
|
|
# compute the previous noisy sample x_t -> x_t-1
|
|
latents, log_prob, _, _ = self.scheduler.step(
|
|
noise_pred,
|
|
t,
|
|
latents,
|
|
generator=generator,
|
|
noise_level=cur_noise_level,
|
|
sde_type=sde_type,
|
|
return_logprobs=logprobs,
|
|
return_dict=False,
|
|
)
|
|
|
|
if i >= sde_window[0] and i < sde_window[1]:
|
|
all_latents.append(latents)
|
|
all_log_probs.append(log_prob)
|
|
all_timesteps.append(t)
|
|
|
|
all_latents = torch.stack(all_latents, dim=1)
|
|
|
|
if all_log_probs[0] is not None:
|
|
all_log_probs = torch.stack(all_log_probs, dim=1)
|
|
else:
|
|
all_log_probs = None
|
|
|
|
all_timesteps = torch.stack(all_timesteps).unsqueeze(0).expand(latents.shape[0], -1)
|
|
|
|
return latents, all_latents, all_log_probs, all_timesteps
|
|
|
|
def forward(
|
|
self,
|
|
req: DiffusionRequestBatch,
|
|
prompt_ids: torch.Tensor | list[int] | None = None,
|
|
prompt_mask: torch.Tensor | None = None,
|
|
negative_prompt_ids: torch.Tensor | list[int] | None = None,
|
|
negative_prompt_mask: torch.Tensor | 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:
|
|
if req.num_reqs != 1:
|
|
raise ValueError("QwenImagePipelineWithLogProbForTest supports only single-request forward.")
|
|
request = req.requests[0]
|
|
|
|
# Extract prompt data from OmniCustomPrompt in req.prompt
|
|
custom_prompt = request.prompt if request.prompt is not None else {}
|
|
if isinstance(custom_prompt, dict):
|
|
prompt_ids = custom_prompt.get("prompt_ids", prompt_ids)
|
|
prompt_mask = custom_prompt.get("prompt_mask", prompt_mask)
|
|
negative_prompt_ids = custom_prompt.get("negative_prompt_ids", negative_prompt_ids)
|
|
negative_prompt_mask = custom_prompt.get("negative_prompt_mask", negative_prompt_mask)
|
|
|
|
# Read sampling params from req.sampling_params
|
|
sp = request.sampling_params
|
|
height = sp.height or self.default_sample_size * self.vae_scale_factor
|
|
width = sp.width or self.default_sample_size * self.vae_scale_factor
|
|
num_inference_steps = sp.num_inference_steps or num_inference_steps
|
|
max_sequence_length = sp.max_sequence_length or max_sequence_length
|
|
|
|
noise_level = sp.extra_args.get("noise_level", None) or noise_level
|
|
sde_window_size = sp.extra_args.get("sde_window_size", None) or sde_window_size
|
|
sde_window_range = sp.extra_args.get("sde_window_range", None) or sde_window_range
|
|
sde_type = sp.extra_args.get("sde_type", None) or sde_type
|
|
logprobs = sp.extra_args.get("logprobs", None) or logprobs
|
|
|
|
generator = sp.generator or generator
|
|
if generator is None and sp.seed is not None:
|
|
generator = torch.Generator(device=self.device).manual_seed(sp.seed)
|
|
true_cfg_scale = sp.true_cfg_scale or true_cfg_scale
|
|
req_num_outputs = getattr(sp, "num_outputs_per_prompt", None)
|
|
if req_num_outputs and req_num_outputs > 0:
|
|
num_images_per_prompt = req_num_outputs
|
|
|
|
self._guidance_scale = guidance_scale
|
|
self._attention_kwargs = attention_kwargs
|
|
self._current_timestep = None
|
|
self._interrupt = False
|
|
|
|
if prompt_ids is not None:
|
|
if isinstance(prompt_ids, list):
|
|
prompt_ids = torch.tensor(prompt_ids, device=self.device)
|
|
batch_size = prompt_ids.shape[0] if prompt_ids.ndim == 2 else 1
|
|
elif prompt_embeds is not None:
|
|
batch_size = prompt_embeds.shape[0]
|
|
else:
|
|
# Both prompt_ids and prompt_embeds are None (e.g. during warmup/dummy run).
|
|
# Return a minimal dummy output to avoid crashing.
|
|
return DiffusionOutput(output=None)
|
|
|
|
if isinstance(negative_prompt_ids, list):
|
|
negative_prompt_ids = torch.tensor(negative_prompt_ids, device=self.device)
|
|
|
|
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_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 = [[(1, height // self.vae_scale_factor // 2, width // self.vae_scale_factor // 2)]] * batch_size
|
|
|
|
timesteps, num_inference_steps = self.prepare_timesteps(num_inference_steps, sigmas, latents.shape[1])
|
|
# num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
|
|
self._num_timesteps = len(timesteps)
|
|
|
|
# handle guidance
|
|
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
|
|
)
|
|
|
|
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()
|
|
end = start + sde_window_size
|
|
sde_window = (start, end)
|
|
else:
|
|
sde_window = (0, len(timesteps) - 1)
|
|
|
|
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]
|
|
|
|
return DiffusionOutput(
|
|
output={
|
|
"payload": {
|
|
"image": _maybe_to_cpu(image),
|
|
},
|
|
"metadata": {
|
|
"prompt_embeddings": {
|
|
"prompt_embeds": _maybe_to_cpu(prompt_embeds),
|
|
"prompt_embeds_mask": _maybe_to_cpu(prompt_embeds_mask),
|
|
"negative_prompt_embeds": _maybe_to_cpu(negative_prompt_embeds),
|
|
"negative_prompt_embeds_mask": _maybe_to_cpu(negative_prompt_embeds_mask),
|
|
},
|
|
},
|
|
},
|
|
trajectory_latents=_maybe_to_cpu(all_latents),
|
|
trajectory_log_probs=_maybe_to_cpu(all_log_probs),
|
|
trajectory_timesteps=_maybe_to_cpu(all_timesteps),
|
|
)
|