Repository: huggingface/diffusers
File: src/diffusers/modular_pipelines/qwenimage/before_denoise.py (Python)

Code before:
        block_state.txt_seq_lens = (
            block_state.prompt_embeds_mask.sum(dim=1).tolist() if block_state.prompt_embeds_mask is not None else None
        )
        block_state.negative_txt_seq_lens = (
            block_state.negative_prompt_embeds_mask.sum(dim=1).tolist()
            if block_state.negative_prompt_embeds_mask is not None
            else None
        )

        self.set_block_state(state, block_state)

        return components, state


class QwenImageEditPlusRoPEInputsStep(QwenImageEditRoPEInputsStep):
    model_name = "qwenimage"
    # TODO: Is there a better way to handle this name? It's used in
    # `QwenImageEditPlusResizeDynamicStep` as well. We can later
    # keep these things as a module-level constant.
    _image_size_output_name = "image_sizes"

    @property
    def inputs(self) -> List[InputParam]:
        inputs_list = super().inputs
        return inputs_list + [
            InputParam(name=self._image_size_output_name, required=True),
        ]

    def __call__(self, components: QwenImageModularPipeline, state: PipelineState) -> PipelineState:
        block_state = self.get_block_state(state)
        vae_image_sizes = getattr(block_state, self._image_size_output_name)
        height, width = block_state.image_height, block_state.image_width

        # for edit, image size can be different from the target size (height/width)
        block_state.img_shapes = [
            [
                (1, height // components.vae_scale_factor // 2, width // components.vae_scale_factor // 2),
                *[
                    (1, vae_height // components.vae_scale_factor // 2, vae_width // components.vae_scale_factor // 2)
                    for vae_width, vae_height in vae_image_sizes
                ],
            ]
        ] * block_state.batch_size

        block_state.txt_seq_lens = (
            block_state.prompt_embeds_mask.sum(dim=1).tolist() if block_state.prompt_embeds_mask is not None else None
        )
        block_state.negative_txt_seq_lens = (
            block_state.negative_prompt_embeds_mask.sum(dim=1).tolist()
            if block_state.negative_prompt_embeds_mask is not None
            else None

Review comment (suggestion):
thanks!, this part is just for this model and when doing the denoise step alone and passing the image latents, it gives an error since the `images_sizes` input is missing.

`ValueError: Required input 'image_sizes' is missing`

Code after:
        block_state.txt_seq_lens = (
            block_state.prompt_embeds_mask.sum(dim=1).tolist() if block_state.prompt_embeds_mask is not None else None
        )
        block_state.negative_txt_seq_lens = (
            block_state.negative_prompt_embeds_mask.sum(dim=1).tolist()
            if block_state.negative_prompt_embeds_mask is not None
            else None
        )

        self.set_block_state(state, block_state)

        return components, state


## ControlNet inputs for denoiser
class QwenImageControlNetBeforeDenoiserStep(ModularPipelineBlocks):
    model_name = "qwenimage"

    @property
    def expected_components(self) -> List[ComponentSpec]:
        return [
            ComponentSpec("controlnet", QwenImageControlNetModel),
        ]

    @property
    def description(self) -> str:
        return "step that prepare inputs for controlnet. Insert before the Denoise Step, after set_timesteps step."

    @property
    def inputs(self) -> List[InputParam]:
        return [
            InputParam("control_guidance_start", default=0.0),
            InputParam("control_guidance_end", default=1.0),
            InputParam("controlnet_conditioning_scale", default=1.0),
            InputParam("control_image_latents", required=True),
            InputParam(
                "timesteps",
                required=True,
                type_hint=torch.Tensor,
                description="The timesteps to use for the denoising process. Can be generated in set_timesteps step.",
            ),
        ]

    @property
    def intermediate_outputs(self) -> List[OutputParam]:
        return [
            OutputParam("controlnet_keep", type_hint=List[float], description="The controlnet keep values"),
        ]

    @torch.no_grad()
    def __call__(self, components: QwenImageModularPipeline, state: PipelineState) -> PipelineState: