diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index be1602788..573cf5ee8 100644 --- a/docs/user_guide/diffusion_features.md +++ b/docs/user_guide/diffusion_features.md @@ -108,7 +108,7 @@ The following tables show which models support each feature: | Model | ⚑TeaCache | ⚑Cache-DiT | πŸ”€SP (Ulysses & Ring) | πŸ”€CFG-Parallel | πŸ”€Tensor-Parallel | πŸ”€HSDP | πŸ’ΎCPU Offload (Layerwise) | πŸ’ΎVAE-Patch-Parallel | πŸ’ΎQuantization | πŸ”„Step Execution | |-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|:----------------:| | **Bagel** | βœ… | βœ… | βœ… | βœ… | βœ… | ❌ | βœ… | ❌ | ❌ | ❌ | -| **FLUX.1-dev** | ❌ | βœ… | ❌ | βœ… | βœ… | βœ… | ❌ | ❌ | βœ… | ❌ | +| **FLUX.1-dev** | βœ… | βœ… | ❌ | βœ… | βœ… | βœ… | ❌ | ❌ | βœ… | ❌ | | **FLUX.1-schnell** | ❌ | βœ… | ❌ | βœ… | βœ… | βœ… | ❌ | ❌ | βœ… | ❌ | | **FLUX.2-klein** | βœ… | βœ… | βœ… | βœ… | βœ… | βœ… | ❌ | ❌ | βœ… | ❌ | | **FLUX.1-Kontext-dev** | ❌ | ❌ | ❌ | ❌ | βœ… | βœ… | ❌ | ❌ | ❌ | ❌ | diff --git a/tests/diffusion/cache/test_teacache_extractors.py b/tests/diffusion/cache/test_teacache_extractors.py index 4bb958a36..b4600dba9 100644 --- a/tests/diffusion/cache/test_teacache_extractors.py +++ b/tests/diffusion/cache/test_teacache_extractors.py @@ -13,6 +13,8 @@ used by TeaCache. Each model's extractor can be tested by: Currently implemented: - TestFlux2KleinExtractor: Flux2Klein model extractor +- TestFlux2Extractor: Flux2 model extractor +- TestFluxExtractor: Flux model extractor """ from abc import ABC, abstractmethod @@ -22,7 +24,12 @@ import pytest import torch from tests.helpers.mark import hardware_test -from vllm_omni.diffusion.cache.teacache.extractors import extract_flux2_context, extract_flux2_klein_context +from vllm_omni.diffusion.cache.teacache.extractors import ( + extract_flux2_context, + extract_flux2_klein_context, + extract_flux_context, +) +from vllm_omni.diffusion.models.flux.flux_transformer import FluxTransformer2DModel from vllm_omni.diffusion.models.flux2_klein.flux2_klein_transformer import ( Flux2Transformer2DModel, ) @@ -234,7 +241,7 @@ class TestFlux2Extractor(BaseExtractorTest): in_channels (128) to inner_dim (6144), so modulated_input should match the projected shape, not the input shape. """ - context = extract_flux2_klein_context(flux2_module, **sample_inputs) + context = extract_flux2_context(flux2_module, **sample_inputs) batch_size, img_seq_len, _ = sample_inputs["hidden_states"].shape inner_dim = flux2_module.inner_dim @@ -277,3 +284,139 @@ class TestFlux2Extractor(BaseExtractorTest): img_ids=torch.randint(0, 64, (1, 1024, 4)), txt_ids=torch.randint(0, 64, (1, 512, 4)), ) + + +@pytest.mark.cpu +class TestFluxExtractor(BaseExtractorTest): + """Test extract_flux_context function.""" + + @pytest.fixture(autouse=True) + def cpu_vllm_config(self): + """Force CPU custom-op dispatch for this test class.""" + from vllm.config import DeviceConfig, VllmConfig, set_current_vllm_config + + with set_current_vllm_config(VllmConfig(device_config=DeviceConfig(device="cpu"))): + yield + + @pytest.fixture(autouse=True) + def mock_flux_attention_backend(self): + """Use the SDPA backend so FLUX can be instantiated in CPU tests.""" + from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend + + with patch("vllm_omni.diffusion.attention.layer.get_attn_backend", return_value=SDPABackend): + yield + + def get_extractor(self): + return extract_flux_context + + @pytest.fixture + def flux_module(self): + """Create a minimal FluxTransformer2DModel for testing.""" + return FluxTransformer2DModel( + num_layers=2, + num_single_layers=2, + num_attention_heads=2, + attention_head_dim=16, + joint_attention_dim=32, + pooled_projection_dim=16, + axes_dims_rope=(4, 4, 8), + ) + + @pytest.fixture + def flux_module_without_guidance(self): + """Create a minimal non-guidance-distilled FLUX transformer.""" + return FluxTransformer2DModel( + num_layers=2, + num_single_layers=2, + num_attention_heads=2, + attention_head_dim=16, + joint_attention_dim=32, + pooled_projection_dim=16, + guidance_embeds=False, + axes_dims_rope=(4, 4, 8), + ) + + def get_module(self, flux_module): + return flux_module + + @pytest.fixture + def sample_inputs(self): + """Create sample input tensors for Flux.""" + batch_size = 1 + img_seq_len = 16 + txt_seq_len = 8 + in_channels = 64 # Flux default in_channels + txt_dim = 32 + pooled_dim = 16 + + return { + "hidden_states": torch.randn(batch_size, img_seq_len, in_channels), + "encoder_hidden_states": torch.randn(batch_size, txt_seq_len, txt_dim), + "pooled_projections": torch.randn(batch_size, pooled_dim), + "timestep": torch.tensor([500]), + "img_ids": torch.randint(0, 64, (batch_size, img_seq_len, 3)), + "txt_ids": torch.randint(0, 64, (batch_size, txt_seq_len, 3)), + "guidance": torch.tensor([3.5]), + } + + def get_sample_inputs(self, sample_inputs): + return sample_inputs + + def test_modulated_input_shape(self, flux_module, sample_inputs): + """Test that modulated_input has the projected FLUX inner dimension.""" + context = extract_flux_context(flux_module, **sample_inputs) + + batch_size, img_seq_len, _ = sample_inputs["hidden_states"].shape + assert context.modulated_input.shape == (batch_size, img_seq_len, flux_module.inner_dim) + + def test_run_transformer_blocks_callable(self, flux_module, sample_inputs): + """Test that run_transformer_blocks is callable.""" + context = extract_flux_context(flux_module, **sample_inputs) + assert callable(context.run_transformer_blocks) + + def test_postprocess_callable(self, flux_module, sample_inputs): + """Test that postprocess is callable.""" + context = extract_flux_context(flux_module, **sample_inputs) + assert callable(context.postprocess) + + def test_postprocess_output_shape(self, flux_module, sample_inputs): + """Test that postprocess projects back to the input channel width.""" + context = extract_flux_context(flux_module, **sample_inputs) + output = context.postprocess(context.hidden_states) + + assert output.sample.shape == sample_inputs["hidden_states"].shape + + def test_postprocess_return_tuple_when_return_dict_false(self, flux_module, sample_inputs): + """Test that postprocess honors return_dict=False.""" + context = extract_flux_context(flux_module, **sample_inputs, return_dict=False) + output = context.postprocess(context.hidden_states) + + assert isinstance(output, tuple) + assert len(output) == 1 + assert output[0].shape == sample_inputs["hidden_states"].shape + + def test_without_guidance(self, flux_module_without_guidance, sample_inputs): + """Test context extraction works for FLUX variants without guidance embeddings.""" + inputs = sample_inputs.copy() + inputs["guidance"] = None + + context = extract_flux_context(flux_module_without_guidance, **inputs) + + assert context is not None + assert context.temb is not None + + def test_invalid_module_raises_error(self): + """Test that invalid module without transformer_blocks raises ValueError.""" + invalid_module = Mock() + invalid_module.transformer_blocks = [] + + with pytest.raises(ValueError, match="Module must have transformer_blocks"): + extract_flux_context( + invalid_module, + hidden_states=torch.randn(1, 16, 64), + encoder_hidden_states=torch.randn(1, 8, 32), + pooled_projections=torch.randn(1, 16), + timestep=torch.tensor([500]), + img_ids=torch.randint(0, 64, (1, 16, 3)), + txt_ids=torch.randint(0, 64, (1, 8, 3)), + ) diff --git a/vllm_omni/diffusion/cache/teacache/extractors.py b/vllm_omni/diffusion/cache/teacache/extractors.py index d0da0d9df..6c5acdbf3 100644 --- a/vllm_omni/diffusion/cache/teacache/extractors.py +++ b/vllm_omni/diffusion/cache/teacache/extractors.py @@ -1069,6 +1069,130 @@ def extract_flux2_context( ) +def extract_flux_context( + module: nn.Module, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + pooled_projections: torch.Tensor, + timestep: torch.LongTensor, + img_ids: torch.Tensor, + txt_ids: torch.Tensor, + guidance: torch.Tensor | None = None, + joint_attention_kwargs: dict[str, Any] | None = None, + **kwargs: Any, +) -> CacheContext: + """ + Extract cache context for FluxTransformer2DModel. + + This mirrors the standard FLUX.1 transformer forward path while exposing + the first modulated image stream tensor as TeaCache's similarity signal. + + Args: + module: FluxTransformer2DModel instance + hidden_states: Input image hidden states tensor + encoder_hidden_states: Text encoder outputs + pooled_projections: Pooled text conditioning + timestep: Current diffusion timestep + img_ids: Image position IDs for RoPE + txt_ids: Text position IDs for RoPE + guidance: Optional guidance scale for guidance-distilled variants + joint_attention_kwargs: Additional attention kwargs + **kwargs: Additional keyword arguments + + Returns: + CacheContext with all information needed for generic caching + """ + from diffusers.models.modeling_outputs import Transformer2DModelOutput + + if not hasattr(module, "transformer_blocks") or len(module.transformer_blocks) == 0: + raise ValueError("Module must have transformer_blocks") + + # ============================================================================ + # PREPROCESSING (Flux1-specific) + # ============================================================================ + hidden_states = module.x_embedder(hidden_states) + timestep = timestep.to(device=hidden_states.device, dtype=hidden_states.dtype) * 1000 + + if guidance is not None: + guidance = guidance.to(device=hidden_states.device, dtype=hidden_states.dtype) * 1000 + + temb = ( + module.time_text_embed(timestep, pooled_projections) + if guidance is None + else module.time_text_embed(timestep, guidance, pooled_projections) + ) + encoder_hidden_states = module.context_embedder(encoder_hidden_states) + + if txt_ids.ndim == 3: + txt_ids = txt_ids[0] + if img_ids.ndim == 3: + img_ids = img_ids[0] + + ids = torch.cat((txt_ids, img_ids), dim=0) + if current_omni_platform.is_npu(): + freqs_cos, freqs_sin = module.pos_embed(ids.cpu()) + image_rotary_emb = (freqs_cos.npu(), freqs_sin.npu()) + else: + image_rotary_emb = module.pos_embed(ids) + + # ============================================================================ + # EXTRACT MODULATED INPUT (for cache decision) + # ============================================================================ + first_block = module.transformer_blocks[0] + modulated_input, *_ = first_block.norm1(hidden_states, emb=temb) + + # ============================================================================ + # DEFINE TRANSFORMER EXECUTION (Flux1-specific) + # ============================================================================ + def run_transformer_blocks() -> tuple[torch.Tensor, torch.Tensor]: + h = hidden_states + e = encoder_hidden_states + + for block in module.transformer_blocks: + e, h = block( + hidden_states=h, + encoder_hidden_states=e, + temb=temb, + image_rotary_emb=image_rotary_emb, + joint_attention_kwargs=joint_attention_kwargs, + ) + + for block in module.single_transformer_blocks: + e, h = block( + hidden_states=h, + encoder_hidden_states=e, + temb=temb, + image_rotary_emb=image_rotary_emb, + joint_attention_kwargs=joint_attention_kwargs, + ) + + return (h, e) + + return_dict = kwargs.get("return_dict", True) + + # ============================================================================ + # DEFINE POSTPROCESSING + # ============================================================================ + def postprocess(h: torch.Tensor) -> Any: + h = module.norm_out(h, temb) + output = module.proj_out(h) + if not return_dict: + return (output,) + return Transformer2DModelOutput(sample=output) + + # ============================================================================ + # RETURN CONTEXT + # ============================================================================ + return CacheContext( + modulated_input=modulated_input, + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + temb=temb, + run_transformer_blocks=run_transformer_blocks, + postprocess=postprocess, + ) + + # Registry for model-specific extractors # Key: Transformer class name # Value: extractor function with signature (module, *args, **kwargs) -> CacheContext @@ -1083,8 +1207,8 @@ EXTRACTOR_REGISTRY: dict[str, Callable] = { "StableAudioDiTModel": extract_stable_audio_context, "Flux2Transformer2DModel": extract_flux2_context, "LongCatImageTransformer2DModel": extract_longcat_context, + "FluxTransformer2DModel": extract_flux_context, # Future models: - # "FluxTransformer2DModel": extract_flux_context, # "CogVideoXTransformer3DModel": extract_cogvideox_context, }