diff --git a/tests/diffusion/models/bagel/test_trajectory_recording.py b/tests/diffusion/models/bagel/test_trajectory_recording.py index 7518388d2..80b3f9d9b 100644 --- a/tests/diffusion/models/bagel/test_trajectory_recording.py +++ b/tests/diffusion/models/bagel/test_trajectory_recording.py @@ -24,15 +24,15 @@ EXPECTED_STEPS = NUM_TIMESTEPS - 1 def _make_mock_bagel(): - """Create a mock Bagel with _forward_flow returning constant velocity.""" + """Create a mock Bagel with forward returning constant velocity.""" mock = MagicMock(spec=Bagel) mock._sp_size = 1 - # _forward_flow returns a small constant velocity so x_t changes each step - def fake_forward_flow(self, x_t, **kwargs): + # forward returns a small constant velocity so x_t changes each step + def fake_forward(self, x_t, **kwargs): return torch.ones_like(x_t) * 0.1 - mock._forward_flow = types.MethodType(fake_forward_flow, mock) + mock.forward = types.MethodType(fake_forward, mock) # _merge_naive_caches is called in the batched CFG path mock._merge_naive_caches = types.MethodType(lambda self, caches: NaiveCache(1), mock) diff --git a/vllm_omni/diffusion/cache/teacache/backend.py b/vllm_omni/diffusion/cache/teacache/backend.py index a5087fe0c..772dec789 100644 --- a/vllm_omni/diffusion/cache/teacache/backend.py +++ b/vllm_omni/diffusion/cache/teacache/backend.py @@ -48,16 +48,7 @@ def enable_bagel_teacache(pipeline: Any, config: DiffusionCacheConfig) -> None: coefficients=config.coefficients, ) transformer = pipeline.bagel - original_forward_flow = transformer._forward_flow - - import types - - def forward_alias(self, *args, **kwargs): - return original_forward_flow(*args, **kwargs) - - transformer.forward = types.MethodType(forward_alias, transformer) apply_teacache_hook(transformer, teacache_config) - transformer._forward_flow = transformer.forward pipeline.transformer = transformer logger.info( diff --git a/vllm_omni/diffusion/cache/teacache/coefficient_estimator.py b/vllm_omni/diffusion/cache/teacache/coefficient_estimator.py index f3a278b21..5dd80718d 100644 --- a/vllm_omni/diffusion/cache/teacache/coefficient_estimator.py +++ b/vllm_omni/diffusion/cache/teacache/coefficient_estimator.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import types from typing import Any import numpy as np @@ -74,15 +73,8 @@ class BagelAdapter: @staticmethod def install_hook(transformer: Any, hook: DataCollectionHook) -> None: - original_forward_flow = transformer._forward_flow - - def forward_alias(self, *args, **kwargs): - return original_forward_flow(*args, **kwargs) - - transformer.forward = types.MethodType(forward_alias, transformer) registry = HookRegistry.get_or_create(transformer) registry.register_hook(hook._HOOK_NAME, hook) - transformer._forward_flow = transformer.forward class StableAudioAdapter: diff --git a/vllm_omni/diffusion/models/bagel/bagel_transformer.py b/vllm_omni/diffusion/models/bagel/bagel_transformer.py index a04ded376..f84807756 100644 --- a/vllm_omni/diffusion/models/bagel/bagel_transformer.py +++ b/vllm_omni/diffusion/models/bagel/bagel_transformer.py @@ -1734,7 +1734,7 @@ class Bagel(nn.Module): packed_seqlens=packed_seqlens, ) - v_t = self._forward_flow_single_branch( + v_t = self.forward_single_branch( **common, packed_indexes=packed_indexes, packed_position_ids=packed_position_ids, @@ -1744,7 +1744,7 @@ class Bagel(nn.Module): ) if cfg_text_scale_ > 1.0: - cfg_text_v_t = self._forward_flow_single_branch( + cfg_text_v_t = self.forward_single_branch( **common, packed_indexes=cfg_text_packed_query_indexes, packed_position_ids=cfg_text_packed_position_ids, @@ -1754,7 +1754,7 @@ class Bagel(nn.Module): ) cfg_img_v_t = None if cfg_img_scale_ > 1.0: - cfg_img_v_t = self._forward_flow_single_branch( + cfg_img_v_t = self.forward_single_branch( **common, packed_indexes=cfg_img_packed_query_indexes, packed_position_ids=cfg_img_packed_position_ids, @@ -1790,7 +1790,7 @@ class Bagel(nn.Module): if use_sp: for i, t in enumerate(timesteps): timestep = torch.tensor([t] * x_t.shape[0], device=x_t.device) - v_t = self._forward_flow_single_branch( + v_t = self.forward_single_branch( x_t=x_t, timestep=timestep, packed_vae_token_indexes=packed_vae_token_indexes, @@ -1883,7 +1883,7 @@ class Bagel(nn.Module): else: cfg_text_scale_ = 1.0 cfg_img_scale_ = 1.0 - v_t = self._forward_flow( + v_t = self.forward( x_t=x_t, timestep=timestep, packed_vae_token_indexes=packed_vae_token_indexes, @@ -2019,7 +2019,7 @@ class Bagel(nn.Module): if use_cfg_this_step: # CFG interval: each rank computes its own branch - local_v_t = self._forward_flow_single_branch( + local_v_t = self.forward_single_branch( x_t=x_t, timestep=timestep, packed_vae_token_indexes=packed_vae_token_indexes, @@ -2046,7 +2046,7 @@ class Bagel(nn.Module): ) else: # Outside CFG interval: all ranks compute with gen inputs, no comm - v_t = self._forward_flow_single_branch( + v_t = self.forward_single_branch( x_t=x_t, timestep=timestep, packed_vae_token_indexes=packed_vae_token_indexes, @@ -2128,7 +2128,7 @@ class Bagel(nn.Module): return v_t - def _forward_flow_single_branch( + def forward_single_branch( self, x_t: torch.Tensor, timestep: torch.LongTensor, @@ -2258,7 +2258,7 @@ class Bagel(nn.Module): v_t = v_t[packed_vae_token_indexes] return v_t - def _forward_flow( + def forward( self, x_t: torch.Tensor, timestep: torch.LongTensor, diff --git a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py index 2c72d9890..13d0cc209 100644 --- a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py +++ b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py @@ -159,8 +159,8 @@ class BagelPipeline(nn.Module, DiffusionPipelineProfilerMixin): self.od_config = od_config self.device = get_local_device() - self._scheduler: object | None = None - self._scheduler_kwargs: dict = {} + self.scheduler: object | None = None + self.scheduler_kwargs: dict = {} model = od_config.model local_files_only = os.path.exists(model) @@ -654,8 +654,8 @@ class BagelPipeline(nn.Module, DiffusionPipelineProfilerMixin): 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"], return_trajectory_latents=req.sampling_params.return_trajectory_latents, - scheduler=self._scheduler, - scheduler_kwargs=self._scheduler_kwargs, + scheduler=self.scheduler, + scheduler_kwargs=self.scheduler_kwargs, ) img = self._decode_image_from_latent(self.bagel, self.vae, latents[0], image_shape)