diff --git a/examples/offline_inference/bagel/end2end.py b/examples/offline_inference/bagel/end2end.py index 0cbe7bb39..bf9968ffa 100644 --- a/examples/offline_inference/bagel/end2end.py +++ b/examples/offline_inference/bagel/end2end.py @@ -46,6 +46,12 @@ def parse_args(): parser.add_argument("--stage-configs-path", type=str, default=None) parser.add_argument("--steps", type=int, default=50, help="Number of inference steps.") + parser.add_argument("--cfg-text-scale", type=float, default=4.0, help="Text CFG scale (default: 4.0)") + parser.add_argument("--cfg-img-scale", type=float, default=1.5, help="Image CFG scale (default: 1.5)") + parser.add_argument( + "--negative-prompt", type=str, default=None, help="Negative prompt (not yet supported, reserved for future)" + ) + args = parser.parse_args() return args @@ -102,6 +108,10 @@ def main(): seed=52, need_kv_receive=False, num_inference_steps=args.steps, + extra_args={ + "cfg_text_scale": args.cfg_text_scale, + "cfg_img_scale": args.cfg_img_scale, + }, ), ) @@ -158,7 +168,12 @@ def main(): if args.modality == "text2img": params_list[0].max_tokens = 1 # type: ignore # The first stage is a SamplingParam (vllm) if len(params_list) > 1: - params_list[1].num_inference_steps = args.steps # type: ignore # The second stage is an OmniDiffusionSamplingParam + diffusion_params = params_list[1] + diffusion_params.num_inference_steps = args.steps # type: ignore + diffusion_params.extra_args = { # type: ignore + "cfg_text_scale": args.cfg_text_scale, + "cfg_img_scale": args.cfg_img_scale, + } omni_outputs = list(omni.generate(prompts=formatted_prompts, sampling_params_list=params_list)) diff --git a/vllm_omni/diffusion/models/bagel/bagel_transformer.py b/vllm_omni/diffusion/models/bagel/bagel_transformer.py index 839fc1d72..7a476ee73 100644 --- a/vllm_omni/diffusion/models/bagel/bagel_transformer.py +++ b/vllm_omni/diffusion/models/bagel/bagel_transformer.py @@ -10,7 +10,6 @@ import math from collections.abc import Iterable from dataclasses import dataclass -from typing import Any import numpy as np import torch @@ -1251,6 +1250,39 @@ class Bagel(nn.Module): def prepare_vae_latent(self, curr_kvlens, curr_rope, image_sizes, new_token_ids): return self.prepare_input(curr_kvlens, curr_rope, image_sizes, new_token_ids) + def prepare_vae_latent_cfg(self, curr_kvlens, curr_rope, image_sizes): + packed_position_ids, packed_indexes, packed_key_value_indexes = list(), list(), list() + + query_curr = curr = 0 + for (H, W), curr_kvlen, curr_position_id in zip(image_sizes, curr_kvlens, curr_rope): + packed_key_value_indexes.extend(range(curr, curr + curr_kvlen)) + curr += curr_kvlen + + packed_indexes.append(curr) + curr += 1 + query_curr += 1 + + h, w = H // self.latent_downsample, W // self.latent_downsample + num_image_tokens = h * w + packed_indexes.extend(range(curr, curr + num_image_tokens)) + curr += num_image_tokens + query_curr += num_image_tokens + + packed_indexes.append(curr) + curr += 1 + query_curr += 1 + + packed_position_ids.extend([curr_position_id] * (num_image_tokens + 2)) + + generation_input = { + "cfg_packed_position_ids": torch.tensor(packed_position_ids, dtype=torch.long), + "cfg_key_values_lens": torch.tensor(curr_kvlens, dtype=torch.int), + "cfg_packed_query_indexes": torch.tensor(packed_indexes, dtype=torch.long), + "cfg_packed_key_value_indexes": torch.tensor(packed_key_value_indexes, dtype=torch.long), + } + + return generation_input + def generate_image( self, packed_text_ids: torch.LongTensor, @@ -1266,11 +1298,25 @@ class Bagel(nn.Module): packed_key_value_indexes: torch.LongTensor, num_timesteps: int = 24, timestep_shift: float = 1.0, + cfg_renorm_min: float = 0.0, + cfg_renorm_type: str = "global", + cfg_interval: tuple[float, float] = [0, 1], + # cfg_text + cfg_text_scale: float = 1.0, + cfg_text_packed_query_indexes: torch.LongTensor | None = None, + cfg_text_packed_position_ids: torch.LongTensor | None = None, + cfg_text_past_key_values: NaiveCache | None = None, + cfg_text_key_values_lens: torch.IntTensor | None = None, + cfg_text_packed_key_value_indexes: torch.LongTensor | None = None, + # cfg_img + cfg_img_scale: float = 1.0, + cfg_img_packed_query_indexes: torch.LongTensor | None = None, + cfg_img_packed_position_ids: torch.LongTensor | None = None, + cfg_img_past_key_values: NaiveCache | None = None, + cfg_img_key_values_lens: torch.IntTensor | None = None, + cfg_img_packed_key_value_indexes: torch.LongTensor | None = None, + cfg_type: str = "parallel", ): - model_pred_cache_dic, model_pred_current = None, None - model_pred_text_cache_dic, model_pred_text_current = None, None - model_pred_img_cache_dic, model_pred_img_current = None, None - x_t = packed_init_noises timesteps = torch.linspace(1, 0, num_timesteps, device=x_t.device) @@ -1280,6 +1326,12 @@ class Bagel(nn.Module): for i, t in enumerate(timesteps): timestep = torch.tensor([t] * x_t.shape[0], device=x_t.device) + if t > cfg_interval[0] and t <= cfg_interval[1]: + cfg_text_scale_ = cfg_text_scale + cfg_img_scale_ = cfg_img_scale + else: + cfg_text_scale_ = 1.0 + cfg_img_scale_ = 1.0 v_t = self._forward_flow( x_t=x_t, timestep=timestep, @@ -1293,17 +1345,26 @@ class Bagel(nn.Module): key_values_lens=key_values_lens, past_key_values=past_key_values, packed_key_value_indexes=packed_key_value_indexes, - # cache - model_pred_cache_dic=model_pred_cache_dic, - model_pred_current=model_pred_current, - model_pred_text_cache_dic=model_pred_text_cache_dic, - model_pred_text_current=model_pred_text_current, - model_pred_img_cache_dic=model_pred_img_cache_dic, - model_pred_img_current=model_pred_img_current, + cfg_renorm_min=cfg_renorm_min, + cfg_renorm_type=cfg_renorm_type, + # cfg_text + cfg_text_scale=cfg_text_scale_, + cfg_text_packed_position_ids=cfg_text_packed_position_ids, + cfg_text_packed_query_indexes=cfg_text_packed_query_indexes, + cfg_text_key_values_lens=cfg_text_key_values_lens, + cfg_text_past_key_values=cfg_text_past_key_values, + cfg_text_packed_key_value_indexes=cfg_text_packed_key_value_indexes, + # cfg_img + cfg_img_scale=cfg_img_scale_, + cfg_img_packed_position_ids=cfg_img_packed_position_ids, + cfg_img_packed_query_indexes=cfg_img_packed_query_indexes, + cfg_img_key_values_lens=cfg_img_key_values_lens, + cfg_img_past_key_values=cfg_img_past_key_values, + cfg_img_packed_key_value_indexes=cfg_img_packed_key_value_indexes, + cfg_type=cfg_type, ) x_t = x_t - v_t.to(x_t.device) * dts[i] # velocity pointing from data to noise - unpacked_latent = x_t.split((packed_seqlens - 2).tolist()) return unpacked_latent @@ -1321,13 +1382,23 @@ class Bagel(nn.Module): key_values_lens: torch.IntTensor, past_key_values: NaiveCache, packed_key_value_indexes: torch.LongTensor, - # cache - model_pred_cache_dic: dict[str, Any] | None = None, - model_pred_current: int | None = None, - model_pred_text_cache_dic: dict[str, Any] | None = None, - model_pred_text_current: int | None = None, - model_pred_img_cache_dic: dict[str, Any] | None = None, - model_pred_img_current: int | None = None, + cfg_renorm_min: float = 0.0, + cfg_renorm_type: str = "global", + # cfg_text + cfg_text_scale: float = 1.0, + cfg_text_packed_position_ids: torch.LongTensor | None = None, + cfg_text_packed_query_indexes: torch.LongTensor | None = None, + cfg_text_key_values_lens: torch.Tensor | None = None, + cfg_text_past_key_values: NaiveCache | None = None, + cfg_text_packed_key_value_indexes: torch.LongTensor | None = None, + # cfg_img + cfg_img_scale: float = 1.0, + cfg_img_packed_position_ids: torch.LongTensor | None = None, + cfg_img_packed_query_indexes: torch.LongTensor | None = None, + cfg_img_key_values_lens: torch.Tensor | None = None, + cfg_img_past_key_values: NaiveCache | None = None, + cfg_img_packed_key_value_indexes: torch.LongTensor | None = None, + cfg_type: str = "parallel", ): packed_text_embedding = self.language_model.model.embed_tokens(packed_text_ids) packed_sequence = packed_text_embedding.new_zeros((sum(packed_seqlens), self.hidden_size)) @@ -1364,4 +1435,70 @@ class Bagel(nn.Module): v_t = self.llm2vae(output.packed_query_sequence) v_t = v_t[packed_vae_token_indexes] + if cfg_text_scale > 1.0: + cfg_text_output = self.language_model.forward( + packed_query_sequence=packed_sequence, + query_lens=packed_seqlens, + packed_query_position_ids=cfg_text_packed_position_ids, + packed_query_indexes=cfg_text_packed_query_indexes, + past_key_values=cfg_text_past_key_values, + key_values_lens=cfg_text_key_values_lens, + packed_key_value_indexes=cfg_text_packed_key_value_indexes, + update_past_key_values=False, + is_causal=False, + **extra_inputs, + ) + cfg_text_v_t = self.llm2vae(cfg_text_output.packed_query_sequence) + cfg_text_v_t = cfg_text_v_t[packed_vae_token_indexes] + + if cfg_img_scale > 1.0: + cfg_img_output = self.language_model.forward( + packed_query_sequence=packed_sequence, + query_lens=packed_seqlens, + packed_query_position_ids=cfg_img_packed_position_ids, + packed_query_indexes=cfg_img_packed_query_indexes, + past_key_values=cfg_img_past_key_values, + key_values_lens=cfg_img_key_values_lens, + packed_key_value_indexes=cfg_img_packed_key_value_indexes, + update_past_key_values=False, + is_causal=False, + **extra_inputs, + ) + cfg_img_v_t = self.llm2vae(cfg_img_output.packed_query_sequence) + cfg_img_v_t = cfg_img_v_t[packed_vae_token_indexes] + + if cfg_text_scale > 1.0: + if cfg_renorm_type == "text_channel": + v_t_text_ = cfg_text_v_t + cfg_text_scale * (v_t - cfg_text_v_t) + norm_v_t = torch.norm(v_t, dim=-1, keepdim=True) + norm_v_t_text_ = torch.norm(v_t_text_, dim=-1, keepdim=True) + scale = (norm_v_t / (norm_v_t_text_ + 1e-8)).clamp(min=cfg_renorm_min, max=1.0) + v_t_text = v_t_text_ * scale + if cfg_img_scale > 1.0: + v_t = cfg_img_v_t + cfg_img_scale * (v_t_text - cfg_img_v_t) + else: + v_t = v_t_text + else: + v_t_text_ = cfg_text_v_t + cfg_text_scale * (v_t - cfg_text_v_t) + + if cfg_img_scale > 1.0: + v_t_ = cfg_img_v_t + cfg_img_scale * (v_t_text_ - cfg_img_v_t) + else: + v_t_ = v_t_text_ + + # NOTE norm is computed over all dimensions, thus currently only supports batch_size = 1 with navit + if cfg_renorm_type == "global": + norm_v_t = torch.norm(v_t) + norm_v_t_ = torch.norm(v_t_) + elif cfg_renorm_type == "channel": + norm_v_t = torch.norm(v_t, dim=-1, keepdim=True) + norm_v_t_ = torch.norm(v_t_, dim=-1, keepdim=True) + else: + raise NotImplementedError(f"{cfg_renorm_type} is not supported") + scale = (norm_v_t / (norm_v_t_ + 1e-8)).clamp(min=cfg_renorm_min, max=1.0) + v_t = v_t_ * scale + else: + # No CFG + pass + return v_t diff --git a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py index 040f7ddec..340188ba7 100644 --- a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py +++ b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py @@ -9,6 +9,7 @@ from __future__ import annotations import json import os from collections.abc import Iterable +from copy import deepcopy from dataclasses import dataclass from math import isqrt @@ -37,7 +38,12 @@ logger = init_logger(__name__) @dataclass class BagelGenParams: num_timesteps: int = 50 - timestep_shift: float = 1.0 + timestep_shift: float = 3.0 + cfg_text_scale: float = 4.0 + cfg_img_scale: float = 1.5 + cfg_interval: tuple = (0.4, 1.0) + cfg_renorm_min: float = 0.0 + cfg_renorm_type: str = "global" def add_special_tokens(tokenizer): @@ -296,10 +302,15 @@ class BagelPipeline(nn.Module): ) image_shape = (height, width) - # Map request params to Bagel gen params (defaults follow Bagel inferencer) + extra_args = getattr(req.sampling_params, "extra_args", {}) or {} + cfg_text_scale = extra_args.get("cfg_text_scale", 4.0) + cfg_img_scale = extra_args.get("cfg_img_scale", 1.5) + gen_params = BagelGenParams( num_timesteps=int(req.sampling_params.num_inference_steps or 50), timestep_shift=3.0, + cfg_text_scale=cfg_text_scale, + cfg_img_scale=cfg_img_scale, ) gen_context = { @@ -307,20 +318,26 @@ class BagelPipeline(nn.Module): "ropes": [0], "past_key_values": NaiveCache(self.bagel.config.llm_config.num_hidden_layers), } + cfg_text_context = deepcopy(gen_context) + cfg_img_context = deepcopy(gen_context) - # Add text prompt (prefill) on gen context. - # [Omni] Check for injected KV Cache from remote transfer injected_kv = req.sampling_params.past_key_values if injected_kv is not None: logger.info("Using injected KV Cache (direct)") gen_context["past_key_values"] = injected_kv - - # User requested: kv_lens and ropes set to [gen_context["past_key_values"].key_cache[0].shape[0]] - # Assuming injected_kv is compatible and has key_cache[0] seq_len = injected_kv.key_cache[0].shape[0] gen_context["kv_lens"] = [seq_len] gen_context["ropes"] = [seq_len] + # Disable CFG: single KV cache cannot support 3-branch CFG + logger.warning("CFG is disabled when using injected KV Cache") + gen_params = BagelGenParams( + num_timesteps=gen_params.num_timesteps, + timestep_shift=gen_params.timestep_shift, + cfg_text_scale=1.0, + cfg_img_scale=1.0, + ) + else: image_input = ( None if isinstance(first_prompt, str) else (first_prompt.get("multi_modal_data") or {}).get("image") @@ -335,9 +352,35 @@ class BagelPipeline(nn.Module): if self.image_processor and self.vae: def vit_transforms(img): - # SigLIP processor returns dict with pixel_values; we want the tensor return self.image_processor(images=img, return_tensors="pt").pixel_values[0] + stride = self.bagel.latent_downsample + max_img_size = int(self.bagel.max_latent_size * stride) + + def _resize_to_stride(img): + if img.mode != "RGB": + img = img.convert("RGB") + w, h = img.size + # Scale down if longest edge exceeds max + scale = min(max_img_size / max(w, h), 1.0) + # Scale up if shortest edge is too small (min 256) + min_img_size = min(256, max_img_size) + scale = max(scale, min_img_size / min(w, h)) + new_w = max(stride, int(round(w * scale / stride) * stride)) + new_h = max(stride, int(round(h * scale / stride) * stride)) + # Clamp to max + new_w = min(new_w, max_img_size) + new_h = min(new_h, max_img_size) + if new_w != w or new_h != h: + img = img.resize((new_w, new_h), Image.BICUBIC) + return img + + image_input = [_resize_to_stride(img) for img in image_input] + + resized_w, resized_h = image_input[0].size + image_shape = (resized_h, resized_w) + logger.info(f"img2img: resized image to {resized_w}x{resized_h}") + def vae_transforms(img): if img.mode != "RGB": img = img.convert("RGB") @@ -345,7 +388,7 @@ class BagelPipeline(nn.Module): arr = torch.from_numpy(np.array(img)).float() / 127.5 - 1.0 return arr.permute(2, 0, 1) - # 1. Update VAE + # Update gen_context with image (VAE + ViT) gen_input_vae, newlens_vae, new_rope_vae = self.bagel.prepare_vae_images( curr_kvlens=gen_context["kv_lens"], curr_rope=gen_context["ropes"], @@ -353,12 +396,9 @@ class BagelPipeline(nn.Module): transforms=vae_transforms, new_token_ids=self.new_token_ids, ) - for k, v in gen_input_vae.items(): if torch.is_tensor(v): gen_input_vae[k] = v.to(self.device) - - # VAE needs bfloat16 to match model strings usually, specifically encode with torch.autocast( device_type=self.device.type, enabled=self.device.type != "cpu", @@ -370,7 +410,6 @@ class BagelPipeline(nn.Module): gen_context["kv_lens"] = newlens_vae gen_context["ropes"] = new_rope_vae - # 2. Update ViT gen_input_img, newlens_img, new_rope_img = self.bagel.prepare_vit_images( curr_kvlens=gen_context["kv_lens"], curr_rope=gen_context["ropes"], @@ -378,11 +417,9 @@ class BagelPipeline(nn.Module): transforms=vit_transforms, new_token_ids=self.new_token_ids, ) - for k, v in gen_input_img.items(): if torch.is_tensor(v): gen_input_img[k] = v.to(self.device) - with torch.autocast( device_type=self.device.type, enabled=self.device.type != "cpu", @@ -393,6 +430,10 @@ class BagelPipeline(nn.Module): ) gen_context["kv_lens"] = newlens_img gen_context["ropes"] = new_rope_img + + cfg_text_context = deepcopy(gen_context) + + # Update gen_context with text prompt generation_input, newlens, new_rope = self.bagel.prepare_prompts( curr_kvlens=gen_context["kv_lens"], curr_rope=gen_context["ropes"], @@ -424,12 +465,56 @@ class BagelPipeline(nn.Module): gen_context["kv_lens"] = newlens gen_context["ropes"] = new_rope + # cfg_text_context: update with negative prompt (no text condition) + neg_prompt = extra_args.get("negative_prompt", "") + neg_input, neg_newlens, neg_rope = self.bagel.prepare_prompts( + curr_kvlens=cfg_text_context["kv_lens"], + curr_rope=cfg_text_context["ropes"], + prompts=[neg_prompt], + tokenizer=self.tokenizer, + new_token_ids=self.new_token_ids, + ) + for k, v in neg_input.items(): + if torch.is_tensor(v): + neg_input[k] = v.to(self.device) + with torch.autocast( + device_type=self.device.type, + enabled=self.device.type != "cpu", + dtype=self.od_config.dtype, + ): + cfg_text_context["past_key_values"] = self.bagel.forward_cache_update_text( + cfg_text_context["past_key_values"], **neg_input + ) + cfg_text_context["kv_lens"] = neg_newlens + cfg_text_context["ropes"] = neg_rope + + # cfg_img_context: update with text prompt (no image condition) + cfg_img_generation_input, cfg_img_newlens, cfg_img_new_rope = self.bagel.prepare_prompts( + curr_kvlens=cfg_img_context["kv_lens"], + curr_rope=cfg_img_context["ropes"], + prompts=[prompt], + tokenizer=self.tokenizer, + new_token_ids=self.new_token_ids, + ) + for k, v in cfg_img_generation_input.items(): + if torch.is_tensor(v): + cfg_img_generation_input[k] = v.to(self.device) + with torch.autocast( + device_type=self.device.type, + enabled=self.device.type != "cpu", + dtype=self.od_config.dtype, + ): + cfg_img_context["past_key_values"] = self.bagel.forward_cache_update_text( + cfg_img_context["past_key_values"], **cfg_img_generation_input + ) + cfg_img_context["kv_lens"] = cfg_img_newlens + cfg_img_context["ropes"] = cfg_img_new_rope + if req.sampling_params.seed is not None: torch.manual_seed(req.sampling_params.seed) if self.device.type == "cuda": torch.cuda.manual_seed(req.sampling_params.seed) - # Prepare latent query and run flow generation_input = self.bagel.prepare_vae_latent( curr_kvlens=gen_context["kv_lens"], curr_rope=gen_context["ropes"], @@ -462,6 +547,25 @@ class BagelPipeline(nn.Module): if torch.is_tensor(v): generation_input[k] = v.to(self.device) + # text cfg + generation_input_cfg_text = self.bagel.prepare_vae_latent_cfg( + curr_kvlens=cfg_text_context["kv_lens"], + curr_rope=cfg_text_context["ropes"], + image_sizes=[image_shape], + ) + # img cfg + generation_input_cfg_img = self.bagel.prepare_vae_latent_cfg( + curr_kvlens=cfg_img_context["kv_lens"], + curr_rope=cfg_img_context["ropes"], + image_sizes=[image_shape], + ) + for k, v in generation_input_cfg_text.items(): + if torch.is_tensor(v): + generation_input_cfg_text[k] = v.to(self.device) + for k, v in generation_input_cfg_img.items(): + if torch.is_tensor(v): + generation_input_cfg_img[k] = v.to(self.device) + with torch.autocast( device_type=self.device.type, enabled=self.device.type != "cpu", @@ -469,12 +573,26 @@ class BagelPipeline(nn.Module): ): latents = self.bagel.generate_image( past_key_values=gen_context["past_key_values"], + cfg_text_past_key_values=cfg_text_context["past_key_values"], + cfg_img_past_key_values=cfg_img_context["past_key_values"], num_timesteps=gen_params.num_timesteps, timestep_shift=gen_params.timestep_shift, + cfg_text_scale=gen_params.cfg_text_scale, + cfg_img_scale=gen_params.cfg_img_scale, + cfg_interval=gen_params.cfg_interval, + cfg_renorm_min=gen_params.cfg_renorm_min, + cfg_renorm_type=gen_params.cfg_renorm_type, **generation_input, + cfg_text_packed_position_ids=generation_input_cfg_text["cfg_packed_position_ids"], + cfg_text_packed_query_indexes=generation_input_cfg_text["cfg_packed_query_indexes"], + cfg_text_key_values_lens=generation_input_cfg_text["cfg_key_values_lens"], + cfg_text_packed_key_value_indexes=generation_input_cfg_text["cfg_packed_key_value_indexes"], + cfg_img_packed_position_ids=generation_input_cfg_img["cfg_packed_position_ids"], + cfg_img_packed_query_indexes=generation_input_cfg_img["cfg_packed_query_indexes"], + cfg_img_key_values_lens=generation_input_cfg_img["cfg_key_values_lens"], + cfg_img_packed_key_value_indexes=generation_input_cfg_img["cfg_packed_key_value_indexes"], ) - # Decode first sample img = self._decode_image_from_latent(self.bagel, self.vae, latents[0], image_shape) return DiffusionOutput(output=img)