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

Code before:
            OutputParam(
                name="txt_seq_lens",
                kwargs_type="denoiser_input_fields",
                type_hint=List[int],
                description="The sequence lengths of the prompt embeds, used for RoPE calculation",
            ),
            OutputParam(
                name="negative_txt_seq_lens",
                kwargs_type="denoiser_input_fields",
                type_hint=List[int],
                description="The sequence lengths of the negative prompt embeds, used for RoPE calculation",
            ),
        ]

    def __call__(self, components: QwenImageModularPipeline, state: PipelineState) -> PipelineState:
        block_state = self.get_block_state(state)

        block_state.img_shapes = [
            [
                (
                    1,
                    block_state.height // components.vae_scale_factor // 2,
                    block_state.width // components.vae_scale_factor // 2,
                )
            ]
            for _ in range(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
        )

        self.set_block_state(state, block_state)

        return components, state


class QwenImageEditRoPEInputsStep(ModularPipelineBlocks):
    model_name = "qwenimage"

    @property
    def description(self) -> str:
        return "Step that prepares the RoPE inputs for denoising process. This is used in QwenImage Edit. Should be placed after prepare_latents step"

    @property
    def inputs(self) -> List[InputParam]:
        return [

Review comment (performance):
We have two options:

1. Either this 
2. Or how it's done in edit: https://github.com/huggingface/diffusers/blob/325a95051bb20787d3db9faba9b4f62b2a63c43a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit.py#L752-L757

Regardless, the current implementation isn't exactly the same as how the standard pipeline implements it and would break for the batched input tests we have.

Code after:
            OutputParam(
                name="txt_seq_lens",
                kwargs_type="denoiser_input_fields",
                type_hint=List[int],
                description="The sequence lengths of the prompt embeds, used for RoPE calculation",
            ),
            OutputParam(
                name="negative_txt_seq_lens",
                kwargs_type="denoiser_input_fields",
                type_hint=List[int],
                description="The sequence lengths of the negative prompt embeds, used for RoPE calculation",
            ),
        ]

    def __call__(self, components: QwenImageModularPipeline, state: PipelineState) -> PipelineState:
        block_state = self.get_block_state(state)

        block_state.img_shapes = [
            [
                (
                    1,
                    block_state.height // components.vae_scale_factor // 2,
                    block_state.width // components.vae_scale_factor // 2,
                )
            ]
        ] * 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
        )

        self.set_block_state(state, block_state)

        return components, state


class QwenImageEditRoPEInputsStep(ModularPipelineBlocks):
    model_name = "qwenimage"

    @property
    def description(self) -> str:
        return "Step that prepares the RoPE inputs for denoising process. This is used in QwenImage Edit. Should be placed after prepare_latents step"

    @property
    def inputs(self) -> List[InputParam]:
        return [
            InputParam(name="batch_size", required=True),