# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import json import math import os import re from collections.abc import Iterable from typing import Any, cast import numpy as np import PIL.Image import torch from diffusers.image_processor import VaeImageProcessor from diffusers.models import AutoencoderKL from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import retrieve_timesteps from diffusers.schedulers import FlowMatchEulerDiscreteScheduler from diffusers.utils.torch_utils import randn_tensor from torch import nn from transformers import ( AutoTokenizer, Qwen2_5_VLForConditionalGeneration, Qwen2VLProcessor, ) from vllm.logger import init_logger 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.utils import get_local_device from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.models.interface import SupportImageInput from vllm_omni.diffusion.models.longcat_image.longcat_image_transformer import ( LongCatImageTransformer2DModel, ) from vllm_omni.diffusion.models.longcat_image.pipeline_longcat_image import calculate_shift from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.inputs.data import OmniTextPrompt from vllm_omni.model_executor.model_loader.weight_utils import ( download_weights_from_hf_specific, ) logger = init_logger(__name__) def get_longcat_image_edit_pre_process_func( od_config: OmniDiffusionConfig, ): """Pre-processing function for LongCatImageEditPipeline.""" 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 = VaeImageProcessor(vae_scale_factor=vae_scale_factor * 2) latent_channels = vae_config.get("latent_channels", 16) def pre_process_func( request: OmniDiffusionRequest, ): """Pre-process requests for LongCatImageEditPipeline.""" for i, prompt in enumerate(request.prompts): multi_modal_data = prompt.get("multi_modal_data", {}) if not isinstance(prompt, str) else None raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None if isinstance(prompt, str): prompt = OmniTextPrompt(prompt=prompt) if "additional_information" not in prompt: prompt["additional_information"] = {} if not raw_image: # None or empty list raise ValueError("""Received no input image. This model requires one input image to run.""") elif isinstance(raw_image, list): if len(raw_image) > 1: raise ValueError( """Received multiple input images. Only a single image is supported by this model.""" ) else: raw_image = raw_image[0] if isinstance(raw_image, str): image = PIL.Image.open(raw_image) else: image = cast(PIL.Image.Image | torch.Tensor | np.ndarray, raw_image) image_size = image.size calculated_width, calculated_height = calculate_dimensions(1024 * 1024, image_size[0] * 1.0 / image_size[1]) height = request.sampling_params.height or calculated_height width = request.sampling_params.width or calculated_width # Store calculated dimensions in request prompt["additional_information"]["calculated_height"] = calculated_height prompt["additional_information"]["calculated_width"] = calculated_width request.sampling_params.height = height request.sampling_params.width = width # Preprocess image if image is not None and not (isinstance(image, torch.Tensor) and image.size(1) == latent_channels): image = image_processor.resize(image, calculated_height, calculated_width) prompt_image = image_processor.resize(image, calculated_height // 2, calculated_width // 2) image = image_processor.preprocess(image, calculated_height, calculated_width) # Store preprocessed image and prompt image in request prompt["additional_information"]["preprocessed_image"] = image prompt["additional_information"]["prompt_image"] = prompt_image request.prompts[i] = prompt return request return pre_process_func def get_longcat_image_post_process_func( od_config: OmniDiffusionConfig, ): 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 = VaeImageProcessor(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.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( encoder_output: torch.Tensor, generator: torch.Generator | None = 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") def calculate_dimensions(target_area, ratio): width = math.sqrt(target_area * ratio) height = width / ratio width = width if width % 16 == 0 else (width // 16 + 1) * 16 height = height if height % 16 == 0 else (height // 16 + 1) * 16 width = int(width) height = int(height) return width, height def prepare_pos_ids(modality_id=0, type="text", start=(0, 0), num_token=None, height=None, width=None): if type == "text": assert num_token if height or width: logger.warning('The parameters of height and width will be ignored in "text" type.') pos_ids = torch.zeros(num_token, 3) pos_ids[..., 0] = modality_id pos_ids[..., 1] = torch.arange(num_token) + start[0] pos_ids[..., 2] = torch.arange(num_token) + start[1] elif type == "image": assert height and width if num_token: logger.warning('The parameter of num_token will be ignored in "image" type.') pos_ids = torch.zeros(height, width, 3) pos_ids[..., 0] = modality_id pos_ids[..., 1] = pos_ids[..., 1] + torch.arange(height)[:, None] + start[0] pos_ids[..., 2] = pos_ids[..., 2] + torch.arange(width)[None, :] + start[1] pos_ids = pos_ids.reshape(height * width, 3) else: raise KeyError(f'Unknown type {type}, only support "text" or "image".') return pos_ids def split_quotation(prompt, quote_pairs=None): """ Implement a regex-based string splitting algorithm that identifies delimiters defined by single or double quote pairs. Examples:: >>> prompt_en = "Please write 'Hello' on the blackboard for me." >>> print(split_quotation(prompt_en)) >>> # output: [('Please write ', False), ("'Hello'", True), (' on the blackboard for me.', False)] """ word_internal_quote_pattern = re.compile(r"[a-zA-Z]+'[a-zA-Z]+") matches_word_internal_quote_pattern = word_internal_quote_pattern.findall(prompt) mapping_word_internal_quote = [] for i, word_src in enumerate(set(matches_word_internal_quote_pattern)): word_tgt = "longcat_$##$_longcat" * (i + 1) prompt = prompt.replace(word_src, word_tgt) mapping_word_internal_quote.append([word_src, word_tgt]) if quote_pairs is None: quote_pairs = [("'", "'"), ('"', '"'), ("‘", "’"), ("“", "”")] pattern = "|".join([re.escape(q1) + r"[^" + re.escape(q1 + q2) + r"]*?" + re.escape(q2) for q1, q2 in quote_pairs]) parts = re.split(f"({pattern})", prompt) result = [] for part in parts: for word_src, word_tgt in mapping_word_internal_quote: part = part.replace(word_tgt, word_src) if re.match(pattern, part): if len(part): result.append((part, True)) else: if len(part): result.append((part, False)) return result class LongCatImageEditPipeline(nn.Module, CFGParallelMixin, SupportImageInput): 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.device = get_local_device() model = od_config.model local_files_only = os.path.exists(model) self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( model, subfolder="scheduler", local_files_only=local_files_only ) self.text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained( model, subfolder="text_encoder", local_files_only=local_files_only ) self.text_processor = Qwen2VLProcessor.from_pretrained( model, subfolder="text_processor", local_files_only=local_files_only ) self.vae = AutoencoderKL.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to( self.device ) self.transformer = LongCatImageTransformer2DModel(od_config=od_config) self.tokenizer = AutoTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only) self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * 2) self.image_processor_vl = self.text_processor.image_processor self.latent_channels = self.vae.config.get("latent_channels", 16) self.image_token = "<|image_pad|>" self.prompt_template_encode_prefix = ( "<|im_start|>system\n" "As an image editing expert, first analyze the content and attributes of the input image(s). " "Then, based on the user's editing instructions, clearly and precisely determine how to modify " "the given image(s), " "ensuring that only the specified parts are altered and all other aspects remain consistent " "with the original(s)." "<|im_end|>\n" "<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>" ) self.prompt_template_encode_suffix = "<|im_end|>\n<|im_start|>assistant\n" self.default_sample_size = 128 self.tokenizer_max_length = 512 def _encode_prompt(self, prompt, image): raw_vl_input = self.image_processor_vl(images=image, return_tensors="pt") pixel_values = raw_vl_input["pixel_values"] image_grid_thw = raw_vl_input["image_grid_thw"] all_tokens = [] for clean_prompt_sub, matched in split_quotation(prompt[0]): if matched: for sub_word in clean_prompt_sub: tokens = self.tokenizer(sub_word, add_special_tokens=False)["input_ids"] all_tokens.extend(tokens) else: tokens = self.tokenizer(clean_prompt_sub, add_special_tokens=False)["input_ids"] all_tokens.extend(tokens) if len(all_tokens) > self.tokenizer_max_length: logger.warning( "Your input was truncated because `max_sequence_length` is set to " f" {self.tokenizer_max_length} input token nums : {len(len(all_tokens))}" ) all_tokens = all_tokens[: self.tokenizer_max_length] text_tokens_and_mask = self.tokenizer.pad( {"input_ids": [all_tokens]}, max_length=self.tokenizer_max_length, padding="max_length", return_attention_mask=True, return_tensors="pt", ) text = self.prompt_template_encode_prefix merge_length = self.image_processor_vl.merge_size**2 while self.image_token in text: num_image_tokens = image_grid_thw.prod() // merge_length text = text.replace(self.image_token, "<|placeholder|>" * num_image_tokens, 1) text = text.replace("<|placeholder|>", self.image_token) prefix_tokens = self.tokenizer(text, add_special_tokens=False)["input_ids"] suffix_tokens = self.tokenizer(self.prompt_template_encode_suffix, add_special_tokens=False)["input_ids"] vision_start_token_id = self.tokenizer.convert_tokens_to_ids("<|vision_start|>") prefix_len = prefix_tokens.index(vision_start_token_id) suffix_len = len(suffix_tokens) prefix_tokens_mask = torch.tensor([1] * len(prefix_tokens), dtype=text_tokens_and_mask.attention_mask[0].dtype) suffix_tokens_mask = torch.tensor([1] * len(suffix_tokens), dtype=text_tokens_and_mask.attention_mask[0].dtype) prefix_tokens = torch.tensor(prefix_tokens, dtype=text_tokens_and_mask.input_ids.dtype) suffix_tokens = torch.tensor(suffix_tokens, dtype=text_tokens_and_mask.input_ids.dtype) input_ids = torch.cat((prefix_tokens, text_tokens_and_mask.input_ids[0], suffix_tokens), dim=-1) attention_mask = torch.cat( (prefix_tokens_mask, text_tokens_and_mask.attention_mask[0], suffix_tokens_mask), dim=-1 ) input_ids = input_ids.unsqueeze(0).to(self.device) attention_mask = attention_mask.unsqueeze(0).to(self.device) pixel_values = pixel_values.to(self.device) image_grid_thw = image_grid_thw.to(self.device) text_output = self.text_encoder( input_ids=input_ids, attention_mask=attention_mask, pixel_values=pixel_values, image_grid_thw=image_grid_thw, output_hidden_states=True, ) # [max_sequence_length, batch, hidden_size] -> [batch, max_sequence_length, hidden_size] # clone to have a contiguous tensor prompt_embeds = text_output.hidden_states[-1].detach() prompt_embeds = prompt_embeds[:, prefix_len:-suffix_len, :] return prompt_embeds def encode_prompt( self, prompt: list[str] = None, image: torch.Tensor | None = None, num_images_per_prompt: int | None = 1, prompt_embeds: torch.Tensor | None = None, ): prompt = [prompt] if isinstance(prompt, str) else prompt batch_size = len(prompt) # If prompt_embeds is provided and prompt is None, skip encoding if prompt_embeds is None: prompt_embeds = self._encode_prompt(prompt, image) _, seq_len, _ = prompt_embeds.shape # duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method 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 = prepare_pos_ids(modality_id=0, type="text", start=(0, 0), num_token=prompt_embeds.shape[1]).to( self.device ) return prompt_embeds, text_ids @staticmethod def _pack_latents(latents, batch_size, num_channels_latents, height, width): latents = latents.view(batch_size, num_channels_latents, height // 2, 2, width // 2, 2) latents = latents.permute(0, 2, 4, 1, 3, 5) latents = latents.reshape(batch_size, (height // 2) * (width // 2), num_channels_latents * 4) return latents @staticmethod def _unpack_latents(latents, height, width, vae_scale_factor): batch_size, num_patches, channels = latents.shape # 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) // (vae_scale_factor * 2)) width = 2 * (int(width) // (vae_scale_factor * 2)) latents = latents.view(batch_size, height // 2, width // 2, channels // 4, 2, 2) latents = latents.permute(0, 3, 1, 4, 2, 5) latents = latents.reshape(batch_size, channels // (2 * 2), height, width) return latents def _encode_vae_image(self, image: torch.Tensor, generator: torch.Generator): if isinstance(generator, list): image_latents = [ retrieve_latents(self.vae.encode(image[i : i + 1]), generator=generator[i], sample_mode="argmax") for i in range(image.shape[0]) ] image_latents = torch.cat(image_latents, dim=0) else: image_latents = retrieve_latents(self.vae.encode(image), generator=generator, sample_mode="argmax") image_latents = (image_latents - self.vae.config.shift_factor) * self.vae.config.scaling_factor return image_latents def prepare_latents( self, image, batch_size, num_channels_latents, height, width, dtype, prompt_embeds_length, device, generator, latents=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)) image_latents, image_latents_ids = None, None if image is not None: image = image.to(device=self.device, dtype=dtype) if image.shape[1] != self.vae.config.latent_channels: image_latents = self._encode_vae_image(image=image, generator=generator) else: image_latents = image if batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] == 0: additional_image_per_prompt = batch_size // image_latents.shape[0] image_latents = torch.cat([image_latents] * additional_image_per_prompt, dim=0) elif batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] != 0: raise ValueError( f"Cannot duplicate `image` of batch size {image_latents.shape[0]} to {batch_size} text prompts." ) else: image_latents = torch.cat([image_latents], dim=0) image_latents = self._pack_latents(image_latents, batch_size, num_channels_latents, height, width) image_latents_ids = prepare_pos_ids( modality_id=2, type="image", start=(prompt_embeds_length, prompt_embeds_length), height=height // 2, width=width // 2, ).to(device, dtype=torch.float64) shape = (batch_size, num_channels_latents, height, width) latents_ids = prepare_pos_ids( modality_id=1, type="image", start=(prompt_embeds_length, prompt_embeds_length), height=height // 2, width=width // 2, ).to(device) 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) latents = self._pack_latents(latents, batch_size, num_channels_latents, height, width) else: latents = latents.to(device=device, dtype=dtype) return latents, image_latents, latents_ids, image_latents_ids def check_inputs( self, prompt, height, width, negative_prompt=None, prompt_embeds=None, negative_prompt_embeds=None ): if height % (self.vae_scale_factor * 2) != 0 or width % (self.vae_scale_factor * 2) != 0: logger.warning( "`height` and `width` have to be divisible by " f"{self.vae_scale_factor * 2} but are {height} and {width}. " "Dimensions will be resized accordingly" ) 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: if isinstance(prompt, str): pass elif isinstance(prompt, list) and len(prompt) == 1: pass else: raise ValueError( f"`prompt` must be a `str` or a `list` of length 1, but is {prompt} (type: {type(prompt)})" ) if negative_prompt is not None and negative_prompt_embeds is not None: raise ValueError( f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:" f" {negative_prompt_embeds}. Please make sure to only forward one of the two." ) def forward( self, req: OmniDiffusionRequest, image: PIL.Image.Image | torch.Tensor | None = None, prompt: str | list[str] | None = None, negative_prompt: str | list[str] | None = None, num_inference_steps: int = 50, sigmas: list[float] | None = None, guidance_scale: float = 3.5, num_images_per_prompt: int | None = 1, generator: torch.Generator | list[torch.Generator] | None = None, latents: torch.FloatTensor | None = None, prompt_embeds: torch.Tensor | None = None, negative_prompt_embeds: torch.Tensor | None = None, output_type: str | None = "pil", return_dict: bool = True, joint_attention_kwargs: dict[str, Any] | None = None, ): # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "") # TODO: May be some data formatting operations on the API side. Hack for now. 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 "") negative_prompt = None if isinstance(first_prompt, str) else first_prompt.get("negative_prompt") prompt_embeds = None if isinstance(first_prompt, str) else first_prompt.get("prompt_embeds") negative_prompt_embeds = None if isinstance(first_prompt, str) else first_prompt.get("negative_prompt_embeds") # type: ignore # Why it is list[torch.Tensor] in OmniTokenInputs or OmniEmbedsPrompt? Doesn't make sense 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 ) num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps num_images_per_prompt = ( req.sampling_params.num_outputs_per_prompt if req.sampling_params.num_outputs_per_prompt is not None else num_images_per_prompt ) generator = req.sampling_params.generator or generator height = req.sampling_params.height or self.default_sample_size * self.vae_scale_factor width = req.sampling_params.width or self.default_sample_size * self.vae_scale_factor if prompt is not None: batch_size = 1 if isinstance(prompt, str) else len(prompt) else: batch_size = prompt_embeds.shape[0] if not isinstance(first_prompt, str) and "preprocessed_image" in ( additional_information := first_prompt.get("additional_information", {}) ): prompt_image = additional_information.get("prompt_image") image = additional_information.get("preprocessed_image") calculated_height = additional_information.get("calculated_height", height) calculated_width = additional_information.get("calculated_width", width) else: image_size = image[0].size if isinstance(image, list) else image.size calculated_width, calculated_height = calculate_dimensions(1024 * 1024, image_size[0] * 1.0 / image_size[1]) if image is not None and not (isinstance(image, torch.Tensor) and image.size(1) == self.latent_channels): image = self.image_processor.resize(image, calculated_height, calculated_width) prompt_image = self.image_processor.resize(image, calculated_height // 2, calculated_width // 2) image = self.image_processor.preprocess(image, calculated_height, calculated_width) self.check_inputs( prompt, calculated_height, calculated_width, negative_prompt=negative_prompt, prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_prompt_embeds, ) (prompt_embeds, text_ids) = self.encode_prompt( prompt=prompt, image=prompt_image, prompt_embeds=prompt_embeds, num_images_per_prompt=num_images_per_prompt ) if guidance_scale > 1: (negative_prompt_embeds, negative_text_ids) = self.encode_prompt( prompt=negative_prompt, image=prompt_image, prompt_embeds=negative_prompt_embeds, num_images_per_prompt=num_images_per_prompt, ) device = self.device # Prepare latent variables num_channels_latents = 16 latents, image_latents, latents_ids, image_latents_ids = self.prepare_latents( image, batch_size * num_images_per_prompt, num_channels_latents, calculated_height, calculated_width, prompt_embeds.dtype, prompt_embeds.shape[1], device, generator, latents, ) # Prepare timesteps sigmas = np.linspace(1.0, 1.0 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas image_seq_len = latents.shape[1] mu = calculate_shift( image_seq_len, self.scheduler.config.get("base_image_seq_len", 256), self.scheduler.config.get("max_image_seq_len", 4096), self.scheduler.config.get("base_shift", 0.5), self.scheduler.config.get("max_shift", 1.15), ) timesteps, num_inference_steps = retrieve_timesteps( self.scheduler, num_inference_steps, device, sigmas=sigmas, mu=mu, ) self._num_timesteps = len(timesteps) guidance = None if image is not None: latent_image_ids = torch.cat([latents_ids, image_latents_ids], dim=0) else: latent_image_ids = latents_ids for i, t in enumerate(timesteps): self._current_timestep = t latent_model_input = latents if image_latents is not None: latent_model_input = torch.cat([latents, image_latents], dim=1) timestep = t.expand(latent_model_input.shape[0]).to(latents.dtype) do_true_cfg = guidance_scale > 1 positive_kwargs = { "hidden_states": latent_model_input, "timestep": timestep / 1000, "guidance": guidance, "encoder_hidden_states": prompt_embeds, "txt_ids": text_ids, "img_ids": latent_image_ids, "return_dict": False, } if do_true_cfg: negative_kwargs = { "hidden_states": latent_model_input, "timestep": timestep / 1000, "encoder_hidden_states": negative_prompt_embeds, "txt_ids": negative_text_ids, "img_ids": latent_image_ids, "return_dict": False, } else: negative_kwargs = None noise_pred = self.predict_noise_maybe_with_cfg( do_true_cfg=do_true_cfg, true_cfg_scale=guidance_scale, positive_kwargs=positive_kwargs, negative_kwargs=negative_kwargs, cfg_normalize=False, output_slice=image_seq_len, ) # 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) self._current_timestep = None if output_type == "latent": image = latents else: latents = self._unpack_latents(latents, calculated_height, calculated_width, self.vae_scale_factor) latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor if latents.dtype != self.vae.dtype: latents = latents.to(dtype=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]: """Load weights using AutoWeightsLoader for vLLM integration.""" loader = AutoWeightsLoader(self) return loader.load_weights(weights)