1142 lines
46 KiB
Python
1142 lines
46 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
import inspect
|
|
import json
|
|
import logging
|
|
import math
|
|
import os
|
|
from collections.abc import Callable, Iterable
|
|
from typing import Any, cast
|
|
|
|
import numpy as np
|
|
import PIL.Image
|
|
import torch
|
|
from diffusers import AutoencoderKLFlux2, FlowMatchEulerDiscreteScheduler
|
|
from diffusers.image_processor import VaeImageProcessor
|
|
from diffusers.pipelines.flux2.pipeline_flux2 import UPSAMPLING_MAX_IMAGE_SIZE
|
|
from diffusers.pipelines.flux2.system_messages import (
|
|
SYSTEM_MESSAGE,
|
|
SYSTEM_MESSAGE_UPSAMPLING_I2I,
|
|
SYSTEM_MESSAGE_UPSAMPLING_T2I,
|
|
)
|
|
from diffusers.utils.torch_utils import randn_tensor
|
|
from torch import nn
|
|
from transformers import AutoConfig, AutoProcessor, PixtralProcessor
|
|
from vllm.model_executor.models.utils import AutoWeightsLoader
|
|
|
|
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
|
|
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
|
|
from vllm_omni.diffusion.distributed.parallel_state import get_classifier_free_guidance_world_size
|
|
from vllm_omni.diffusion.distributed.utils import get_local_device
|
|
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
|
|
from vllm_omni.diffusion.models.flux2 import Flux2Transformer2DModel
|
|
from vllm_omni.diffusion.models.interface import SupportImageInput
|
|
from vllm_omni.diffusion.models.mistral_encoder import MistralEncoderModel
|
|
from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
|
|
from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
|
|
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
|
from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs
|
|
from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class Flux2ImageProcessor(VaeImageProcessor):
|
|
"""Image processor to preprocess the reference image for Flux2."""
|
|
|
|
def __init__(
|
|
self,
|
|
do_resize: bool = True,
|
|
vae_scale_factor: int = 16,
|
|
vae_latent_channels: int = 32,
|
|
do_normalize: bool = True,
|
|
do_convert_rgb: bool = True,
|
|
):
|
|
super().__init__(
|
|
do_resize=do_resize,
|
|
vae_scale_factor=vae_scale_factor,
|
|
vae_latent_channels=vae_latent_channels,
|
|
do_normalize=do_normalize,
|
|
do_convert_rgb=do_convert_rgb,
|
|
)
|
|
|
|
@staticmethod
|
|
def check_image_input(
|
|
image: PIL.Image.Image,
|
|
max_aspect_ratio: int = 8,
|
|
min_side_length: int = 64,
|
|
max_area: int = 1024 * 1024,
|
|
) -> PIL.Image.Image:
|
|
if not isinstance(image, PIL.Image.Image):
|
|
raise ValueError(f"Image must be a PIL.Image.Image, got {type(image)}")
|
|
|
|
width, height = image.size
|
|
if width < min_side_length or height < min_side_length:
|
|
raise ValueError(f"Image too small: {width}x{height}. Both dimensions must be at least {min_side_length}px")
|
|
|
|
aspect_ratio = max(width / height, height / width)
|
|
if aspect_ratio > max_aspect_ratio:
|
|
raise ValueError(
|
|
f"Aspect ratio too extreme: {width}x{height} (ratio: {aspect_ratio:.1f}:1). "
|
|
f"Maximum allowed ratio is {max_aspect_ratio}:1"
|
|
)
|
|
|
|
if width * height > max_area:
|
|
logger.warning("Image area exceeds recommended maximum; resizing will be applied.")
|
|
|
|
return image
|
|
|
|
@staticmethod
|
|
def _resize_to_target_area(image: PIL.Image.Image, target_area: int = 1024 * 1024) -> PIL.Image.Image:
|
|
image_width, image_height = image.size
|
|
scale = math.sqrt(target_area / (image_width * image_height))
|
|
width = int(image_width * scale)
|
|
height = int(image_height * scale)
|
|
return image.resize((width, height), PIL.Image.Resampling.LANCZOS)
|
|
|
|
@staticmethod
|
|
def _resize_if_exceeds_area(image: PIL.Image.Image, target_area: int = 1024 * 1024) -> PIL.Image.Image:
|
|
image_width, image_height = image.size
|
|
if image_width * image_height <= target_area:
|
|
return image
|
|
return Flux2ImageProcessor._resize_to_target_area(image, target_area)
|
|
|
|
def _resize_and_crop(self, image: PIL.Image.Image, width: int, height: int) -> PIL.Image.Image:
|
|
image_width, image_height = image.size
|
|
left = (image_width - width) // 2
|
|
top = (image_height - height) // 2
|
|
right = left + width
|
|
bottom = top + height
|
|
return image.crop((left, top, right, bottom))
|
|
|
|
@staticmethod
|
|
def concatenate_images(images: list[PIL.Image.Image]) -> PIL.Image.Image:
|
|
if len(images) == 1:
|
|
return images[0].copy()
|
|
|
|
images = [img.convert("RGB") if img.mode != "RGB" else img for img in images]
|
|
total_width = sum(img.width for img in images)
|
|
max_height = max(img.height for img in images)
|
|
background_color = (255, 255, 255)
|
|
new_img = PIL.Image.new("RGB", (total_width, max_height), background_color)
|
|
|
|
x_offset = 0
|
|
for img in images:
|
|
y_offset = (max_height - img.height) // 2
|
|
new_img.paste(img, (x_offset, y_offset))
|
|
x_offset += img.width
|
|
|
|
return new_img
|
|
|
|
|
|
def get_flux2_post_process_func(
|
|
od_config: OmniDiffusionConfig,
|
|
):
|
|
if od_config.output_type == "latent":
|
|
return lambda x: x
|
|
model_name = od_config.model
|
|
if os.path.exists(model_name):
|
|
model_path = model_name
|
|
else:
|
|
model_path = download_weights_from_hf_specific(model_name, None, ["*"])
|
|
|
|
vae_config_path = os.path.join(model_path, "vae/config.json")
|
|
with open(vae_config_path) as f:
|
|
vae_config = json.load(f)
|
|
vae_scale_factor = 2 ** (len(vae_config["block_out_channels"]) - 1) if "block_out_channels" in vae_config else 8
|
|
|
|
image_processor = Flux2ImageProcessor(vae_scale_factor=vae_scale_factor * 2)
|
|
|
|
def post_process_func(images: torch.Tensor):
|
|
return image_processor.postprocess(images)
|
|
|
|
return post_process_func
|
|
|
|
|
|
# Copied from diffusers.pipelines.flux2.pipeline_flux2.format_input
|
|
def format_input(
|
|
prompts: list[str],
|
|
system_message: str = SYSTEM_MESSAGE,
|
|
images: list[PIL.Image.Image] | list[list[PIL.Image.Image]] = None,
|
|
) -> list[list[dict[str, Any]]]:
|
|
"""
|
|
Format a batch of text prompts into the conversation format expected by apply_chat_template. Optionally, add images
|
|
to the input.
|
|
|
|
Args:
|
|
prompts: List of text prompts
|
|
system_message: System message to use (default: CREATIVE_SYSTEM_MESSAGE)
|
|
images (optional): List of images to add to the input.
|
|
|
|
Returns:
|
|
`list[list[dict[str, Any]]]`: List of conversations, where each conversation is a list of message dicts
|
|
"""
|
|
# Remove [IMG] tokens from prompts to avoid Pixtral validation issues
|
|
# when truncation is enabled. The processor counts [IMG] tokens and fails
|
|
# if the count changes after truncation.
|
|
cleaned_txt = [prompt.replace("[IMG]", "") for prompt in prompts]
|
|
|
|
if images is None or len(images) == 0:
|
|
return [
|
|
[
|
|
{
|
|
"role": "system",
|
|
"content": [{"type": "text", "text": system_message}],
|
|
},
|
|
{"role": "user", "content": [{"type": "text", "text": prompt}]},
|
|
]
|
|
for prompt in cleaned_txt
|
|
]
|
|
else:
|
|
assert len(images) == len(prompts), "Number of images must match number of prompts"
|
|
messages = [
|
|
[
|
|
{
|
|
"role": "system",
|
|
"content": [{"type": "text", "text": system_message}],
|
|
},
|
|
]
|
|
for _ in cleaned_txt
|
|
]
|
|
|
|
for i, (el, batch_images) in enumerate(zip(messages, images)):
|
|
# optionally add the images per batch element.
|
|
if batch_images is not None:
|
|
el.append(
|
|
{
|
|
"role": "user",
|
|
"content": [{"type": "image", "image": image_obj} for image_obj in batch_images],
|
|
}
|
|
)
|
|
# add the text.
|
|
el.append(
|
|
{
|
|
"role": "user",
|
|
"content": [{"type": "text", "text": cleaned_txt[i]}],
|
|
}
|
|
)
|
|
|
|
return messages
|
|
|
|
|
|
# Copied from diffusers.pipelines.flux2.pipeline_flux2._validate_and_process_images
|
|
def _validate_and_process_images(
|
|
images: list[list[PIL.Image.Image]] | list[PIL.Image.Image],
|
|
image_processor: Flux2ImageProcessor,
|
|
upsampling_max_image_size: int,
|
|
) -> list[list[PIL.Image.Image]]:
|
|
# Simple validation: ensure it's a list of PIL images or list of lists of PIL images
|
|
if not images:
|
|
return []
|
|
|
|
# Check if it's a list of lists or a list of images
|
|
if isinstance(images[0], PIL.Image.Image):
|
|
# It's a list of images, convert to list of lists
|
|
images = [[im] for im in images]
|
|
|
|
# potentially concatenate multiple images to reduce the size
|
|
images = [[image_processor.concatenate_images(img_i)] if len(img_i) > 1 else img_i for img_i in images]
|
|
|
|
# cap the pixels
|
|
images = [
|
|
[image_processor._resize_if_exceeds_area(img_i, upsampling_max_image_size) for img_i in img_i]
|
|
for img_i in images
|
|
]
|
|
return images
|
|
|
|
|
|
# Copied from diffusers.pipelines.flux2.pipeline_flux2.compute_empirical_mu
|
|
def compute_empirical_mu(image_seq_len: int, num_steps: int) -> float:
|
|
a1, b1 = 8.73809524e-05, 1.89833333
|
|
a2, b2 = 0.00016927, 0.45666666
|
|
|
|
if image_seq_len > 4300:
|
|
mu = a2 * image_seq_len + b2
|
|
return float(mu)
|
|
|
|
m_200 = a2 * image_seq_len + b2
|
|
m_10 = a1 * image_seq_len + b1
|
|
|
|
a = (m_200 - m_10) / 190.0
|
|
b = m_200 - 200.0 * a
|
|
mu = a * num_steps + b
|
|
|
|
return float(mu)
|
|
|
|
|
|
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
|
|
def retrieve_timesteps(
|
|
scheduler,
|
|
num_inference_steps: int | None = None,
|
|
device: str | torch.device | None = None,
|
|
timesteps: list[int] | None = None,
|
|
sigmas: list[float] | None = None,
|
|
**kwargs,
|
|
) -> tuple[torch.Tensor, int]:
|
|
r"""
|
|
Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
|
|
custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
|
|
|
|
Args:
|
|
scheduler (`SchedulerMixin`):
|
|
The scheduler to get timesteps from.
|
|
num_inference_steps (`int`):
|
|
The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
|
|
must be `None`.
|
|
device (`str` or `torch.device`, *optional*):
|
|
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
|
|
timesteps (`List[int]`, *optional*):
|
|
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
|
|
`num_inference_steps` and `sigmas` must be `None`.
|
|
sigmas (`List[float]`, *optional*):
|
|
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
|
|
`num_inference_steps` and `timesteps` must be `None`.
|
|
|
|
Returns:
|
|
`tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
|
|
second element is the number of inference steps.
|
|
"""
|
|
if timesteps is not None and sigmas is not None:
|
|
raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
|
|
if timesteps is not None:
|
|
accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
|
|
if not accepts_timesteps:
|
|
raise ValueError(
|
|
f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
|
|
f" timestep schedules. Please check whether you are using the correct scheduler."
|
|
)
|
|
scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
|
|
timesteps = scheduler.timesteps
|
|
num_inference_steps = len(timesteps)
|
|
elif sigmas is not None:
|
|
accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
|
|
if not accept_sigmas:
|
|
raise ValueError(
|
|
f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
|
|
f" sigmas schedules. Please check whether you are using the correct scheduler."
|
|
)
|
|
scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
|
|
timesteps = scheduler.timesteps
|
|
num_inference_steps = len(timesteps)
|
|
else:
|
|
scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
|
|
timesteps = scheduler.timesteps
|
|
return timesteps, num_inference_steps
|
|
|
|
|
|
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
|
|
def retrieve_latents(encoder_output: torch.Tensor, generator: torch.Generator = None, sample_mode: str = "sample"):
|
|
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
|
|
return encoder_output.latent_dist.sample(generator)
|
|
elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax":
|
|
return encoder_output.latent_dist.mode()
|
|
elif hasattr(encoder_output, "latents"):
|
|
return encoder_output.latents
|
|
else:
|
|
raise AttributeError("Could not access latents of provided encoder_output")
|
|
|
|
|
|
class Flux2Pipeline(nn.Module, CFGParallelMixin, SupportImageInput, ProgressBarMixin, DiffusionPipelineProfilerMixin):
|
|
"""Flux2 pipeline for text-to-image generation."""
|
|
|
|
_callback_tensor_inputs = ["latents", "prompt_embeds"]
|
|
|
|
support_image_input = True
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
od_config: OmniDiffusionConfig,
|
|
prefix: str = "",
|
|
):
|
|
super().__init__()
|
|
self.od_config = od_config
|
|
|
|
self.weights_sources = [
|
|
DiffusersPipelineLoader.ComponentSource(
|
|
model_or_path=od_config.model,
|
|
subfolder="transformer",
|
|
revision=None,
|
|
prefix="transformer.",
|
|
fall_back_to_pt=True,
|
|
)
|
|
]
|
|
|
|
self._execution_device = get_local_device()
|
|
model = od_config.model
|
|
# Check if model is a local path
|
|
local_files_only = os.path.exists(model)
|
|
|
|
self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
|
|
model, subfolder="scheduler", local_files_only=local_files_only
|
|
)
|
|
self.weights_sources.append(
|
|
DiffusersPipelineLoader.ComponentSource(
|
|
model_or_path=od_config.model,
|
|
subfolder="text_encoder",
|
|
revision=None,
|
|
prefix="text_encoder.",
|
|
fall_back_to_pt=True,
|
|
),
|
|
)
|
|
text_encoder_config = AutoConfig.from_pretrained(
|
|
model, subfolder="text_encoder", local_files_only=local_files_only
|
|
)
|
|
self.text_encoder = MistralEncoderModel(
|
|
text_encoder_config,
|
|
prefix="text_encoder",
|
|
).to(self._execution_device)
|
|
self.tokenizer = PixtralProcessor.from_pretrained(
|
|
model, subfolder="tokenizer", local_files_only=local_files_only
|
|
)
|
|
self.text_encoder.set_processor(
|
|
self.tokenizer,
|
|
system_message_t2i=SYSTEM_MESSAGE_UPSAMPLING_T2I,
|
|
system_message_i2i=SYSTEM_MESSAGE_UPSAMPLING_I2I,
|
|
)
|
|
self.vae = AutoencoderKLFlux2.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to(
|
|
self._execution_device
|
|
)
|
|
transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, Flux2Transformer2DModel)
|
|
self.transformer = Flux2Transformer2DModel(
|
|
quant_config=od_config.quantization_config, od_config=od_config, **transformer_kwargs
|
|
)
|
|
|
|
self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8
|
|
self.image_processor = Flux2ImageProcessor(vae_scale_factor=self.vae_scale_factor * 2)
|
|
self.tokenizer_max_length = 512
|
|
self.default_sample_size = 128
|
|
|
|
self.system_message = SYSTEM_MESSAGE
|
|
self.system_message_upsampling_t2i = SYSTEM_MESSAGE_UPSAMPLING_T2I
|
|
self.system_message_upsampling_i2i = SYSTEM_MESSAGE_UPSAMPLING_I2I
|
|
self.upsampling_max_image_size = UPSAMPLING_MAX_IMAGE_SIZE
|
|
|
|
self._guidance_scale = None
|
|
self._attention_kwargs = None
|
|
self._num_timesteps = None
|
|
|
|
self.setup_diffusion_pipeline_profiler(
|
|
enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler
|
|
)
|
|
self._current_timestep = None
|
|
self._interrupt = False
|
|
|
|
@staticmethod
|
|
def _get_mistral_3_small_prompt_embeds(
|
|
text_encoder: MistralEncoderModel,
|
|
tokenizer: AutoProcessor,
|
|
prompt: str | list[str],
|
|
dtype: torch.dtype | None = None,
|
|
device: torch.device | None = None,
|
|
max_sequence_length: int = 512,
|
|
system_message: str = SYSTEM_MESSAGE,
|
|
hidden_states_layers: list[int] = (10, 20, 30),
|
|
):
|
|
dtype = text_encoder.dtype if dtype is None else dtype
|
|
device = text_encoder.device if device is None else device
|
|
|
|
prompt = [prompt] if isinstance(prompt, str) else prompt
|
|
|
|
# Format input messages
|
|
messages_batch = format_input(prompts=prompt, system_message=system_message)
|
|
|
|
# Process all messages at once
|
|
inputs = tokenizer.apply_chat_template(
|
|
messages_batch,
|
|
add_generation_prompt=False,
|
|
tokenize=True,
|
|
return_dict=True,
|
|
return_tensors="pt",
|
|
padding="max_length",
|
|
truncation=True,
|
|
max_length=max_sequence_length,
|
|
)
|
|
|
|
# Move to device
|
|
input_ids = inputs["input_ids"].to(device)
|
|
attention_mask = inputs["attention_mask"].to(device)
|
|
|
|
# Forward pass through the model
|
|
output = text_encoder(
|
|
input_ids=input_ids,
|
|
attention_mask=attention_mask,
|
|
output_hidden_states=True,
|
|
use_cache=False,
|
|
)
|
|
|
|
# Only use outputs from intermediate layers and stack them
|
|
out = torch.stack([output.hidden_states[k] for k in hidden_states_layers], dim=1)
|
|
out = out.to(dtype=dtype, device=device)
|
|
|
|
batch_size, num_channels, seq_len, hidden_dim = out.shape
|
|
prompt_embeds = out.permute(0, 2, 1, 3).reshape(batch_size, seq_len, num_channels * hidden_dim)
|
|
|
|
return prompt_embeds
|
|
|
|
@staticmethod
|
|
# Copied from diffusers.pipelines.flux2.pipeline_flux2.Flux2Pipeline._prepare_text_ids
|
|
def _prepare_text_ids(
|
|
x: torch.Tensor, # (B, L, D) or (L, D)
|
|
t_coord: torch.Tensor | None = None,
|
|
):
|
|
B, L, _ = x.shape
|
|
out_ids = []
|
|
|
|
for i in range(B):
|
|
t = torch.arange(1) if t_coord is None else t_coord[i]
|
|
h = torch.arange(1)
|
|
w = torch.arange(1)
|
|
seq_positions = torch.arange(L)
|
|
|
|
coords = torch.cartesian_prod(t, h, w, seq_positions)
|
|
out_ids.append(coords)
|
|
|
|
return torch.stack(out_ids)
|
|
|
|
@staticmethod
|
|
# Copied from diffusers.pipelines.flux2.pipeline_flux2.Flux2Pipeline._prepare_latent_ids
|
|
def _prepare_latent_ids(
|
|
latents: torch.Tensor, # (B, C, H, W)
|
|
):
|
|
r"""
|
|
Generates 4D position coordinates (T, H, W, L) for latent tensors.
|
|
|
|
Args:
|
|
latents (torch.Tensor):
|
|
Latent tensor of shape (B, C, H, W)
|
|
|
|
Returns:
|
|
torch.Tensor:
|
|
Position IDs tensor of shape (B, H*W, 4) All batches share the same coordinate structure: T=0,
|
|
H=[0..H-1], W=[0..W-1], L=0
|
|
"""
|
|
|
|
batch_size, _, height, width = latents.shape
|
|
|
|
t = torch.arange(1) # [0] - time dimension
|
|
h = torch.arange(height)
|
|
w = torch.arange(width)
|
|
layer_ids = torch.arange(1) # [0] - layer dimension
|
|
|
|
# Create position IDs: (H*W, 4)
|
|
latent_ids = torch.cartesian_prod(t, h, w, layer_ids)
|
|
|
|
# Expand to batch: (B, H*W, 4)
|
|
latent_ids = latent_ids.unsqueeze(0).expand(batch_size, -1, -1)
|
|
|
|
return latent_ids
|
|
|
|
@staticmethod
|
|
# Copied from diffusers.pipelines.flux2.pipeline_flux2.Flux2Pipeline._prepare_image_ids
|
|
def _prepare_image_ids(
|
|
image_latents: list[torch.Tensor], # [(1, C, H, W), (1, C, H, W), ...]
|
|
scale: int = 10,
|
|
):
|
|
r"""
|
|
Generates 4D time-space coordinates (T, H, W, L) for a sequence of image latents.
|
|
|
|
This function creates a unique coordinate for every pixel/patch across all input latent with different
|
|
dimensions.
|
|
|
|
Args:
|
|
image_latents (List[torch.Tensor]):
|
|
A list of image latent feature tensors, typically of shape (C, H, W).
|
|
scale (int, optional):
|
|
A factor used to define the time separation (T-coordinate) between latents. T-coordinate for the i-th
|
|
latent is: 'scale + scale * i'. Defaults to 10.
|
|
|
|
Returns:
|
|
torch.Tensor:
|
|
The combined coordinate tensor. Shape: (1, N_total, 4) Where N_total is the sum of (H * W) for all
|
|
input latents.
|
|
|
|
Coordinate Components (Dimension 4):
|
|
- T (Time): The unique index indicating which latent image the coordinate belongs to.
|
|
- H (Height): The row index within that latent image.
|
|
- W (Width): The column index within that latent image.
|
|
- L (Seq. Length): A sequence length dimension, which is always fixed at 0 (size 1)
|
|
"""
|
|
|
|
if not isinstance(image_latents, list):
|
|
raise ValueError(f"Expected `image_latents` to be a list, got {type(image_latents)}.")
|
|
|
|
# create time offset for each reference image
|
|
t_coords = [scale + scale * t for t in torch.arange(0, len(image_latents))]
|
|
t_coords = [t.view(-1) for t in t_coords]
|
|
|
|
image_latent_ids = []
|
|
for x, t in zip(image_latents, t_coords):
|
|
x = x.squeeze(0)
|
|
_, height, width = x.shape
|
|
|
|
x_ids = torch.cartesian_prod(t, torch.arange(height), torch.arange(width), torch.arange(1))
|
|
image_latent_ids.append(x_ids)
|
|
|
|
image_latent_ids = torch.cat(image_latent_ids, dim=0)
|
|
image_latent_ids = image_latent_ids.unsqueeze(0)
|
|
|
|
return image_latent_ids
|
|
|
|
@staticmethod
|
|
# Copied from diffusers.pipelines.flux2.pipeline_flux2.Flux2Pipeline._patchify_latents
|
|
def _patchify_latents(latents):
|
|
batch_size, num_channels_latents, height, width = latents.shape
|
|
latents = latents.view(batch_size, num_channels_latents, height // 2, 2, width // 2, 2)
|
|
latents = latents.permute(0, 1, 3, 5, 2, 4)
|
|
latents = latents.reshape(batch_size, num_channels_latents * 4, height // 2, width // 2)
|
|
return latents
|
|
|
|
@staticmethod
|
|
# Copied from diffusers.pipelines.flux2.pipeline_flux2.Flux2Pipeline._unpatchify_latents
|
|
def _unpatchify_latents(latents):
|
|
batch_size, num_channels_latents, height, width = latents.shape
|
|
latents = latents.reshape(batch_size, num_channels_latents // (2 * 2), 2, 2, height, width)
|
|
latents = latents.permute(0, 1, 4, 2, 5, 3)
|
|
latents = latents.reshape(batch_size, num_channels_latents // (2 * 2), height * 2, width * 2)
|
|
return latents
|
|
|
|
@staticmethod
|
|
# Copied from diffusers.pipelines.flux2.pipeline_flux2.Flux2Pipeline._pack_latents
|
|
def _pack_latents(latents):
|
|
"""
|
|
pack latents: (batch_size, num_channels, height, width) -> (batch_size, height * width, num_channels)
|
|
"""
|
|
|
|
batch_size, num_channels, height, width = latents.shape
|
|
latents = latents.reshape(batch_size, num_channels, height * width).permute(0, 2, 1)
|
|
|
|
return latents
|
|
|
|
@staticmethod
|
|
# Copied from diffusers.pipelines.flux2.pipeline_flux2.Flux2Pipeline._unpack_latents_with_ids
|
|
def _unpack_latents_with_ids(x: torch.Tensor, x_ids: torch.Tensor) -> list[torch.Tensor]:
|
|
"""
|
|
using position ids to scatter tokens into place
|
|
"""
|
|
x_list = []
|
|
for data, pos in zip(x, x_ids):
|
|
_, ch = data.shape # noqa: F841
|
|
h_ids = pos[:, 1].to(torch.int64)
|
|
w_ids = pos[:, 2].to(torch.int64)
|
|
|
|
h = torch.max(h_ids) + 1
|
|
w = torch.max(w_ids) + 1
|
|
|
|
flat_ids = h_ids * w + w_ids
|
|
|
|
out = torch.zeros((h * w, ch), device=data.device, dtype=data.dtype)
|
|
out.scatter_(0, flat_ids.unsqueeze(1).expand(-1, ch), data)
|
|
|
|
# reshape from (H * W, C) to (H, W, C) and permute to (C, H, W)
|
|
|
|
out = out.view(h, w, ch).permute(2, 0, 1)
|
|
x_list.append(out)
|
|
|
|
return torch.stack(x_list, dim=0)
|
|
|
|
def upsample_prompt(
|
|
self,
|
|
prompt: str | list[str],
|
|
images: list[PIL.Image.Image] | list[list[PIL.Image.Image]] = None,
|
|
temperature: float = 0.15,
|
|
device: torch.device = None,
|
|
) -> list[str]:
|
|
if images:
|
|
images = _validate_and_process_images(images, self.image_processor, self.upsampling_max_image_size)
|
|
return self.text_encoder.upsample_prompt(
|
|
prompt,
|
|
images=images,
|
|
temperature=temperature,
|
|
device=device,
|
|
)
|
|
|
|
def encode_prompt(
|
|
self,
|
|
prompt: str | list[str],
|
|
device: torch.device | None = None,
|
|
num_images_per_prompt: int = 1,
|
|
prompt_embeds: torch.Tensor | None = None,
|
|
max_sequence_length: int = 512,
|
|
text_encoder_out_layers: tuple[int, ...] = (10, 20, 30),
|
|
):
|
|
device = device or self._execution_device
|
|
|
|
if prompt is None:
|
|
prompt = ""
|
|
|
|
prompt = [prompt] if isinstance(prompt, str) else prompt
|
|
|
|
if prompt_embeds is None:
|
|
prompt_embeds = self._get_mistral_3_small_prompt_embeds(
|
|
text_encoder=self.text_encoder,
|
|
tokenizer=self.tokenizer,
|
|
prompt=prompt,
|
|
device=device,
|
|
max_sequence_length=max_sequence_length,
|
|
system_message=self.system_message,
|
|
hidden_states_layers=text_encoder_out_layers,
|
|
)
|
|
|
|
batch_size, 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)
|
|
|
|
text_ids = self._prepare_text_ids(prompt_embeds)
|
|
text_ids = text_ids.to(device)
|
|
return prompt_embeds, text_ids
|
|
|
|
# Copied from diffusers.pipelines.flux2.pipeline_flux2.Flux2Pipeline._encode_vae_image
|
|
def _encode_vae_image(self, image: torch.Tensor, generator: torch.Generator):
|
|
if image.ndim != 4:
|
|
raise ValueError(f"Expected image dims 4, got {image.ndim}.")
|
|
|
|
image_latents = retrieve_latents(self.vae.encode(image), generator=generator, sample_mode="argmax")
|
|
image_latents = self._patchify_latents(image_latents)
|
|
|
|
latents_bn_mean = self.vae.bn.running_mean.view(1, -1, 1, 1).to(image_latents.device, image_latents.dtype)
|
|
latents_bn_std = torch.sqrt(self.vae.bn.running_var.view(1, -1, 1, 1) + self.vae.config.batch_norm_eps)
|
|
image_latents = (image_latents - latents_bn_mean) / latents_bn_std
|
|
|
|
return image_latents
|
|
|
|
# Copied from diffusers.pipelines.flux2.pipeline_flux2.Flux2Pipeline.prepare_latents
|
|
def prepare_latents(
|
|
self,
|
|
batch_size,
|
|
num_latents_channels,
|
|
height,
|
|
width,
|
|
dtype,
|
|
device,
|
|
generator: torch.Generator,
|
|
latents: torch.Tensor | None = None,
|
|
):
|
|
# VAE applies 8x compression on images but we must also account for packing which requires
|
|
# latent height and width to be divisible by 2.
|
|
height = 2 * (int(height) // (self.vae_scale_factor * 2))
|
|
width = 2 * (int(width) // (self.vae_scale_factor * 2))
|
|
|
|
shape = (batch_size, num_latents_channels * 4, height // 2, width // 2)
|
|
if isinstance(generator, list) and len(generator) != batch_size:
|
|
raise ValueError(
|
|
f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
|
|
f" size of {batch_size}. Make sure the batch size matches the length of the generators."
|
|
)
|
|
if latents is None:
|
|
latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
|
|
else:
|
|
latents = latents.to(device=device, dtype=dtype)
|
|
|
|
latent_ids = self._prepare_latent_ids(latents)
|
|
latent_ids = latent_ids.to(device)
|
|
|
|
latents = self._pack_latents(latents) # [B, C, H, W] -> [B, H*W, C]
|
|
return latents, latent_ids
|
|
|
|
# Copied from diffusers.pipelines.flux2.pipeline_flux2.Flux2Pipeline.prepare_image_latents
|
|
def prepare_image_latents(
|
|
self,
|
|
images: list[torch.Tensor],
|
|
batch_size,
|
|
generator: torch.Generator,
|
|
device,
|
|
dtype,
|
|
):
|
|
image_latents = []
|
|
for image in images:
|
|
image = image.to(device=device, dtype=dtype)
|
|
image_latent = self._encode_vae_image(image=image, generator=generator)
|
|
image_latents.append(image_latent) # (1, 128, 32, 32)
|
|
|
|
image_latent_ids = self._prepare_image_ids(image_latents)
|
|
|
|
# Pack each latent and concatenate
|
|
packed_latents = []
|
|
for latent in image_latents:
|
|
# latent: (1, 128, 32, 32)
|
|
packed = self._pack_latents(latent) # (1, 1024, 128)
|
|
packed = packed.squeeze(0) # (1024, 128) - remove batch dim
|
|
packed_latents.append(packed)
|
|
|
|
# Concatenate all reference tokens along sequence dimension
|
|
image_latents = torch.cat(packed_latents, dim=0) # (N*1024, 128)
|
|
image_latents = image_latents.unsqueeze(0) # (1, N*1024, 128)
|
|
|
|
image_latents = image_latents.repeat(batch_size, 1, 1)
|
|
image_latent_ids = image_latent_ids.repeat(batch_size, 1, 1)
|
|
image_latent_ids = image_latent_ids.to(device)
|
|
|
|
return image_latents, image_latent_ids
|
|
|
|
def check_inputs(
|
|
self,
|
|
prompt,
|
|
height,
|
|
width,
|
|
prompt_embeds=None,
|
|
callback_on_step_end_tensor_inputs=None,
|
|
):
|
|
if (
|
|
height is not None
|
|
and height % (self.vae_scale_factor * 2) != 0
|
|
or width is not None
|
|
and width % (self.vae_scale_factor * 2) != 0
|
|
):
|
|
logger.warning(
|
|
"`height` and `width` have to be divisible by %s but are %s and %s. "
|
|
"Dimensions will be resized accordingly",
|
|
self.vae_scale_factor * 2,
|
|
height,
|
|
width,
|
|
)
|
|
|
|
if callback_on_step_end_tensor_inputs is not None and not all(
|
|
k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs
|
|
):
|
|
raise ValueError(
|
|
f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, "
|
|
f"but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
|
|
)
|
|
|
|
if prompt is not None and prompt_embeds is not None:
|
|
raise ValueError(
|
|
f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
|
|
" only forward one of the two."
|
|
)
|
|
elif prompt is None and prompt_embeds is None:
|
|
raise ValueError(
|
|
"Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
|
|
)
|
|
elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
|
|
raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
|
|
|
|
@property
|
|
def guidance_scale(self):
|
|
return self._guidance_scale
|
|
|
|
@property
|
|
def attention_kwargs(self):
|
|
return self._attention_kwargs
|
|
|
|
@property
|
|
def num_timesteps(self):
|
|
return self._num_timesteps
|
|
|
|
@property
|
|
def current_timestep(self):
|
|
return self._current_timestep
|
|
|
|
@property
|
|
def interrupt(self):
|
|
return self._interrupt
|
|
|
|
def check_cfg_parallel_validity(self, true_cfg_scale: float, has_neg_prompt: bool):
|
|
if get_classifier_free_guidance_world_size() == 1:
|
|
return True
|
|
|
|
if true_cfg_scale <= 1:
|
|
logger.warning("CFG parallel is NOT working correctly when true_cfg_scale <= 1.")
|
|
return False
|
|
|
|
if not has_neg_prompt:
|
|
logger.warning(
|
|
"CFG parallel is NOT working correctly when there is no negative prompt or negative prompt embeddings."
|
|
)
|
|
return False
|
|
return True
|
|
|
|
def forward(
|
|
self,
|
|
req: OmniDiffusionRequest,
|
|
image: PIL.Image.Image | list[PIL.Image.Image] | None = None,
|
|
prompt: str | list[str] | None = None,
|
|
height: int | None = None,
|
|
width: int | None = None,
|
|
num_inference_steps: int = 50,
|
|
sigmas: list[float] | None = None,
|
|
guidance_scale: float | None = 4.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,
|
|
negative_prompt_embeds: torch.Tensor | None = None,
|
|
output_type: str | None = "pil",
|
|
return_dict: bool = True,
|
|
attention_kwargs: dict[str, Any] | None = None,
|
|
callback_on_step_end: Callable[[int, int, dict], None] | None = None,
|
|
callback_on_step_end_tensor_inputs: list[str] = ["latents"],
|
|
max_sequence_length: int = 512,
|
|
text_encoder_out_layers: tuple[int, ...] = (10, 20, 30),
|
|
caption_upsample_temperature: float = None,
|
|
) -> DiffusionOutput:
|
|
if len(req.prompts) > 1:
|
|
logger.warning(
|
|
"""This model only supports a single prompt, not a batched request.""",
|
|
"""Taking only the first image for now.""",
|
|
)
|
|
first_prompt = req.prompts[0]
|
|
prompt = first_prompt if isinstance(first_prompt, str) else (first_prompt.get("prompt") or "")
|
|
|
|
if (
|
|
raw_image := None
|
|
if isinstance(first_prompt, str)
|
|
else first_prompt.get("multi_modal_data", {}).get("image")
|
|
) is None:
|
|
pass # use image from param list
|
|
elif isinstance(raw_image, list):
|
|
image = [PIL.Image.open(im) if isinstance(im, str) else cast(PIL.Image.Image, im) for im in raw_image]
|
|
else:
|
|
image = PIL.Image.open(raw_image) if isinstance(raw_image, str) else cast(PIL.Image.Image, raw_image)
|
|
|
|
height = req.sampling_params.height or height
|
|
width = req.sampling_params.width or width
|
|
num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps
|
|
sigmas = req.sampling_params.sigmas or sigmas
|
|
guidance_scale = (
|
|
req.sampling_params.guidance_scale if req.sampling_params.guidance_scale is not None else guidance_scale
|
|
)
|
|
generator = req.sampling_params.generator or generator
|
|
num_images_per_prompt = (
|
|
req.sampling_params.num_outputs_per_prompt
|
|
if req.sampling_params.num_outputs_per_prompt > 0
|
|
else num_images_per_prompt
|
|
)
|
|
max_sequence_length = req.sampling_params.max_sequence_length or max_sequence_length
|
|
text_encoder_out_layers = req.sampling_params.extra_args.get("text_encoder_out_layers", text_encoder_out_layers)
|
|
caption_upsample_temperature = req.sampling_params.extra_args.get(
|
|
"caption_upsample_temperature", caption_upsample_temperature
|
|
)
|
|
|
|
req_prompt_embeds = [p.get("prompt_embeds") if not isinstance(p, str) else None for p in req.prompts]
|
|
if any(p is not None for p in req_prompt_embeds):
|
|
# If at list one prompt is provided as an embedding,
|
|
# Then assume that the user wants to provide embeddings for all prompts, and enter this if block
|
|
# If the user in fact provides mixed input format, req_prompt_embeds will have some None's
|
|
# And `torch.stack` automatically raises an exception for us
|
|
prompt_embeds = torch.stack(req_prompt_embeds) # type: ignore # intentionally expect TypeError
|
|
|
|
req_negative_prompt_embeds = [
|
|
p.get("negative_prompt_embeds") if not isinstance(p, str) else None for p in req.prompts
|
|
]
|
|
if all(p is not None for p in req_negative_prompt_embeds):
|
|
negative_prompt_embeds = torch.stack(req_negative_prompt_embeds) # type: ignore # intentionally expect TypeError
|
|
|
|
req_negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in req.prompts]
|
|
|
|
# 1. Check inputs. Raise error if not correct
|
|
self.check_inputs(
|
|
prompt=prompt,
|
|
height=height,
|
|
width=width,
|
|
prompt_embeds=prompt_embeds,
|
|
callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
|
|
)
|
|
|
|
self._guidance_scale = guidance_scale
|
|
self._attention_kwargs = attention_kwargs
|
|
self._current_timestep = None
|
|
self._interrupt = False
|
|
guidance_tensor = None
|
|
|
|
# 2. Define call parameters
|
|
if prompt is not None and isinstance(prompt, str):
|
|
batch_size = 1
|
|
elif prompt is not None and isinstance(prompt, list):
|
|
batch_size = len(prompt)
|
|
else:
|
|
batch_size = prompt_embeds.shape[0]
|
|
|
|
device = self._execution_device
|
|
|
|
# 3. prepare text embeddings
|
|
if caption_upsample_temperature:
|
|
prompt = self.upsample_prompt(prompt, images=image, temperature=caption_upsample_temperature, device=device)
|
|
prompt_embeds, text_ids = self.encode_prompt(
|
|
prompt=prompt,
|
|
prompt_embeds=prompt_embeds,
|
|
device=device,
|
|
num_images_per_prompt=num_images_per_prompt,
|
|
max_sequence_length=max_sequence_length,
|
|
text_encoder_out_layers=text_encoder_out_layers,
|
|
)
|
|
|
|
has_neg_prompt = negative_prompt_embeds is not None or any(req_negative_prompt)
|
|
do_true_cfg = self.guidance_scale > 1 and has_neg_prompt
|
|
|
|
self.check_cfg_parallel_validity(self.guidance_scale, has_neg_prompt)
|
|
negative_text_ids = None
|
|
if do_true_cfg:
|
|
negative_prompt = req_negative_prompt
|
|
negative_prompt_embeds, negative_text_ids = self.encode_prompt(
|
|
prompt=negative_prompt,
|
|
prompt_embeds=negative_prompt_embeds,
|
|
device=device,
|
|
num_images_per_prompt=num_images_per_prompt,
|
|
max_sequence_length=max_sequence_length,
|
|
text_encoder_out_layers=text_encoder_out_layers,
|
|
)
|
|
|
|
# 4. process images
|
|
if image is not None and not isinstance(image, list):
|
|
image = [image]
|
|
|
|
condition_images = None
|
|
if image is not None:
|
|
for img in image:
|
|
self.image_processor.check_image_input(img)
|
|
|
|
condition_images = []
|
|
for img in image:
|
|
image_width, image_height = img.size
|
|
if image_width * image_height > 1024 * 1024:
|
|
img = self.image_processor._resize_to_target_area(img, 1024 * 1024)
|
|
image_width, image_height = img.size
|
|
|
|
multiple_of = self.vae_scale_factor * 2
|
|
image_width = (image_width // multiple_of) * multiple_of
|
|
image_height = (image_height // multiple_of) * multiple_of
|
|
img = self.image_processor.preprocess(img, height=image_height, width=image_width, resize_mode="crop")
|
|
condition_images.append(img)
|
|
height = height or image_height
|
|
width = width or image_width
|
|
|
|
height = height or self.default_sample_size * self.vae_scale_factor
|
|
width = width or self.default_sample_size * self.vae_scale_factor
|
|
|
|
# 5. prepare latent variables
|
|
num_channels_latents = self.transformer.config.in_channels // 4
|
|
latents, latent_ids = self.prepare_latents(
|
|
batch_size=batch_size * num_images_per_prompt,
|
|
num_latents_channels=num_channels_latents,
|
|
height=height,
|
|
width=width,
|
|
dtype=prompt_embeds.dtype,
|
|
device=device,
|
|
generator=generator,
|
|
latents=latents,
|
|
)
|
|
|
|
image_latents = None
|
|
image_latent_ids = None
|
|
if condition_images is not None:
|
|
image_latents, image_latent_ids = self.prepare_image_latents(
|
|
images=condition_images,
|
|
batch_size=batch_size * num_images_per_prompt,
|
|
generator=generator,
|
|
device=device,
|
|
dtype=self.vae.dtype,
|
|
)
|
|
|
|
# 6. Prepare timesteps
|
|
sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas
|
|
if hasattr(self.scheduler.config, "use_flow_sigmas") and self.scheduler.config.use_flow_sigmas:
|
|
sigmas = None
|
|
image_seq_len = latents.shape[1]
|
|
mu = compute_empirical_mu(image_seq_len=image_seq_len, num_steps=num_inference_steps)
|
|
timesteps, num_inference_steps = retrieve_timesteps(
|
|
self.scheduler,
|
|
num_inference_steps,
|
|
device,
|
|
sigmas=sigmas,
|
|
mu=mu,
|
|
)
|
|
self._num_timesteps = len(timesteps)
|
|
|
|
# handle guidance
|
|
if self.transformer.guidance_embeds is not None:
|
|
guidance_tensor = torch.full([1], self.guidance_scale, device=device, dtype=torch.float32)
|
|
guidance_tensor = guidance_tensor.expand(latents.shape[0])
|
|
|
|
# For editing pipelines, we need to slice the output to remove condition latents
|
|
output_slice = latents.size(1) if image_latents is not None else None
|
|
|
|
# 7. Denoising loop
|
|
# We set the index here to remove DtoH sync, helpful especially during compilation.
|
|
# Check out more details here: https://github.com/huggingface/diffusers/pull/11696
|
|
self.scheduler.set_begin_index(0)
|
|
with self.progress_bar(total=len(timesteps)) as pbar:
|
|
for i, t in enumerate(timesteps):
|
|
if self.interrupt:
|
|
continue
|
|
|
|
self._current_timestep = t
|
|
timestep = t.expand(latents.shape[0]).to(latents.dtype)
|
|
|
|
latent_model_input = latents.to(self.transformer.dtype)
|
|
latent_image_ids = latent_ids
|
|
|
|
if image_latents is not None:
|
|
latent_model_input = torch.cat([latents, image_latents], dim=1).to(self.transformer.dtype)
|
|
latent_image_ids = torch.cat([latent_ids, image_latent_ids], dim=1)
|
|
|
|
positive_kwargs = {
|
|
"hidden_states": latent_model_input,
|
|
"timestep": timestep / 1000,
|
|
"guidance": guidance_tensor,
|
|
"encoder_hidden_states": prompt_embeds,
|
|
"txt_ids": text_ids,
|
|
"img_ids": latent_image_ids,
|
|
"joint_attention_kwargs": self.attention_kwargs,
|
|
"return_dict": False,
|
|
}
|
|
if do_true_cfg:
|
|
negative_kwargs = {
|
|
"hidden_states": latent_model_input,
|
|
"timestep": timestep / 1000,
|
|
"guidance": guidance_tensor,
|
|
"encoder_hidden_states": negative_prompt_embeds,
|
|
"txt_ids": negative_text_ids,
|
|
"img_ids": latent_image_ids,
|
|
"joint_attention_kwargs": self.attention_kwargs,
|
|
"return_dict": False,
|
|
}
|
|
else:
|
|
negative_kwargs = None
|
|
|
|
noise_pred = self.predict_noise_maybe_with_cfg(
|
|
do_true_cfg=do_true_cfg,
|
|
true_cfg_scale=self.guidance_scale,
|
|
positive_kwargs=positive_kwargs,
|
|
negative_kwargs=negative_kwargs,
|
|
cfg_normalize=False,
|
|
output_slice=output_slice,
|
|
)
|
|
|
|
# Compute the previous noisy sample x_t -> x_t-1 with automatic CFG sync
|
|
latents = self.scheduler_step_maybe_with_cfg(noise_pred, t, latents, do_true_cfg)
|
|
|
|
if callback_on_step_end is not None:
|
|
callback_kwargs = {}
|
|
for k in callback_on_step_end_tensor_inputs:
|
|
callback_kwargs[k] = locals()[k]
|
|
callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
|
|
|
|
latents = callback_outputs.pop("latents", latents)
|
|
prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
|
|
|
|
pbar.update()
|
|
|
|
self._current_timestep = None
|
|
|
|
latents = self._unpack_latents_with_ids(latents, latent_ids)
|
|
|
|
latents_bn_mean = self.vae.bn.running_mean.view(1, -1, 1, 1).to(latents.device, latents.dtype)
|
|
latents_bn_std = torch.sqrt(self.vae.bn.running_var.view(1, -1, 1, 1) + self.vae.config.batch_norm_eps).to(
|
|
latents.device, latents.dtype
|
|
)
|
|
latents = latents * latents_bn_std + latents_bn_mean
|
|
latents = self._unpatchify_latents(latents)
|
|
if output_type == "latent":
|
|
image = latents
|
|
else:
|
|
if latents.dtype != self.vae.dtype:
|
|
latents = latents.to(self.vae.dtype)
|
|
image = self.vae.decode(latents, return_dict=False)[0]
|
|
|
|
return DiffusionOutput(output=image)
|
|
|
|
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
|
loader = AutoWeightsLoader(self)
|
|
return loader.load_weights(weights)
|