qwen modular refactor, unpack before decode

Author: yiyixuxu

src/diffusers/modular_pipelines/qwenimage/decoders.py

@@ -29,6 +29,45 @@
 
 logger = logging.get_logger(__name__)
 
+class QwenImageAfterDenoiseStep(ModularPipelineBlocks):
+    model_name = "qwenimage"
+
+    @property
+    def description(self) -> str:
+        return "Step that unpack the latents from 3D tensor (batch_size, sequence_length, channels) into 5D tensor (batch_size, channels, 1, height, width)"
+
+    @property
+    def expected_components(self) -> List[ComponentSpec]:
+        components = [
+            ComponentSpec("pachifier", QwenImagePachifier, default_creation_method="from_config"),
+        ]
+
+        return components
+
+    @property
+    def inputs(self) -> List[InputParam]:
+        return [
+            InputParam(name="height", required=True),
+            InputParam(name="width", required=True),
+            InputParam(
+                name="latents",
+                required=True,
+                type_hint=torch.Tensor,
+                description="The latents to decode, can be generated in the denoise step",
+            ),
+        ]
+
+    @torch.no_grad()
+    def __call__(self, components: QwenImageModularPipeline, state: PipelineState) -> PipelineState:
+        block_state = self.get_block_state(state)
+
+        vae_scale_factor = components.vae_scale_factor
+        block_state.latents = components.pachifier.unpack_latents(
+            block_state.latents, block_state.height, block_state.width, vae_scale_factor=vae_scale_factor
+        )
+
+        self.set_block_state(state, block_state)
+        return components, state
 
 class QwenImageDecoderStep(ModularPipelineBlocks):
     model_name = "qwenimage"
@@ -41,16 +80,13 @@ def description(self) -> str:
     def expected_components(self) -> List[ComponentSpec]:
         components = [
             ComponentSpec("vae", AutoencoderKLQwenImage),
-            ComponentSpec("pachifier", QwenImagePachifier, default_creation_method="from_config"),
         ]
 
         return components
 
     @property
     def inputs(self) -> List[InputParam]:
         return [
-            InputParam(name="height", required=True),
-            InputParam(name="width", required=True),
             InputParam(
                 name="latents",
                 required=True,
@@ -74,10 +110,10 @@ def __call__(self, components: QwenImageModularPipeline, state: PipelineState) -
         block_state = self.get_block_state(state)
 
         # YiYi Notes: remove support for output_type = "latents', we can just skip decode/encode step in modular
-        vae_scale_factor = components.vae_scale_factor
-        block_state.latents = components.pachifier.unpack_latents(
-            block_state.latents, block_state.height, block_state.width, vae_scale_factor=vae_scale_factor
-        )
+        if block_state.latents.ndim == 4:
+            block_state.latents = block_state.latents.unsqueeze(dim=1)
+        elif block_state.latents.ndim != 5:
+            raise ValueError(f"expect latents to be a 4D or 5D tensor but got: {block_state.latents.shape}. Please make sure the latents are unpacked before decode step.")
         block_state.latents = block_state.latents.to(components.vae.dtype)
 
         latents_mean = (

src/diffusers/modular_pipelines/qwenimage/modular_blocks.py

@@ -26,7 +26,7 @@
     QwenImageSetTimestepsStep,
     QwenImageSetTimestepsWithStrengthStep,
 )
-from .decoders import QwenImageDecoderStep, QwenImageInpaintProcessImagesOutputStep, QwenImageProcessImagesOutputStep
+from .decoders import QwenImageAfterDenoiseStep, QwenImageDecoderStep, QwenImageInpaintProcessImagesOutputStep, QwenImageProcessImagesOutputStep
 from .denoise import (
     QwenImageControlNetDenoiseStep,
     QwenImageDenoiseStep,
@@ -92,6 +92,7 @@ def description(self):
         ("set_timesteps", QwenImageSetTimestepsStep()),
         ("prepare_rope_inputs", QwenImageRoPEInputsStep()),
         ("denoise", QwenImageDenoiseStep()),
+        ("after_denoise", QwenImageAfterDenoiseStep()),
         ("decode", QwenImageDecodeStep()),
     ]
 )
@@ -205,6 +206,7 @@ def description(self):
         ("prepare_inpaint_latents", QwenImageInpaintPrepareLatentsStep()),
         ("prepare_rope_inputs", QwenImageRoPEInputsStep()),
         ("denoise", QwenImageInpaintDenoiseStep()),
+        ("after_denoise", QwenImageAfterDenoiseStep()),
         ("decode", QwenImageInpaintDecodeStep()),
     ]
 )
@@ -264,6 +266,7 @@ def description(self):
         ("prepare_img2img_latents", QwenImagePrepareLatentsWithStrengthStep()),
         ("prepare_rope_inputs", QwenImageRoPEInputsStep()),
         ("denoise", QwenImageDenoiseStep()),
+        ("after_denoise", QwenImageAfterDenoiseStep()),
         ("decode", QwenImageDecodeStep()),
     ]
 )
@@ -529,8 +532,9 @@ class QwenImageCoreDenoiseStep(SequentialPipelineBlocks):
         QwenImageAutoBeforeDenoiseStep,
         QwenImageOptionalControlNetBeforeDenoiseStep,
         QwenImageAutoDenoiseStep,
+        QwenImageAfterDenoiseStep,
     ]
-    block_names = ["input", "controlnet_input", "before_denoise", "controlnet_before_denoise", "denoise"]
+    block_names = ["input", "controlnet_input", "before_denoise", "controlnet_before_denoise", "denoise", "after_denoise"]
 
     @property
     def description(self):
@@ -653,6 +657,7 @@ def description(self):
         ("set_timesteps", QwenImageSetTimestepsStep()),
         ("prepare_rope_inputs", QwenImageEditRoPEInputsStep()),
         ("denoise", QwenImageEditDenoiseStep()),
+        ("after_denoise", QwenImageAfterDenoiseStep()),
         ("decode", QwenImageDecodeStep()),
     ]
 )
@@ -702,6 +707,7 @@ def description(self) -> str:
         ("prepare_inpaint_latents", QwenImageInpaintPrepareLatentsStep()),
         ("prepare_rope_inputs", QwenImageEditRoPEInputsStep()),
         ("denoise", QwenImageEditInpaintDenoiseStep()),
+        ("after_denoise", QwenImageAfterDenoiseStep()),
         ("decode", QwenImageInpaintDecodeStep()),
     ]
 )
@@ -841,8 +847,9 @@ class QwenImageEditCoreDenoiseStep(SequentialPipelineBlocks):
         QwenImageEditAutoInputStep,
         QwenImageEditAutoBeforeDenoiseStep,
         QwenImageEditAutoDenoiseStep,
+        QwenImageAfterDenoiseStep,
     ]
-    block_names = ["input", "before_denoise", "denoise"]
+    block_names = ["input", "before_denoise", "denoise", "after_denoise"]
 
     @property
     def description(self):
@@ -954,6 +961,7 @@ class QwenImageEditPlusInputStep(SequentialPipelineBlocks):
         ("set_timesteps", QwenImageSetTimestepsStep()),
         ("prepare_rope_inputs", QwenImageEditPlusRoPEInputsStep()),
         ("denoise", QwenImageEditDenoiseStep()),
+        ("after_denoise", QwenImageAfterDenoiseStep()),
         ("decode", QwenImageDecodeStep()),
     ]
 )
@@ -1037,8 +1045,9 @@ class QwenImageEditPlusCoreDenoiseStep(SequentialPipelineBlocks):
         QwenImageEditPlusAutoInputStep,
         QwenImageEditPlusAutoBeforeDenoiseStep,
         QwenImageEditAutoDenoiseStep,
+        QwenImageAfterDenoiseStep,
     ]
-    block_names = ["input", "before_denoise", "denoise"]
+    block_names = ["input", "before_denoise", "denoise", "after_denoise"]
 
     @property
     def description(self):