Repository: Project-MONAI/MONAI
File: monai/apps/generation/maisi/networks/controlnet_maisi.py (Python)

Code before:
            with_conditioning,
            transformer_num_layers,
            cross_attention_dim,
            num_class_embeds,
            upcast_attention,
            use_flash_attention,
            conditioning_embedding_in_channels,
            conditioning_embedding_num_channels,
        )


    def forward(
        self,
        x: torch.Tensor,
        timesteps: torch.Tensor,
        controlnet_cond: torch.Tensor,
        conditioning_scale: float = 1.0,
        context: torch.Tensor | None = None,
        class_labels: torch.Tensor | None = None,
    ) -> tuple[tuple[torch.Tensor], torch.Tensor]:
        """
        Args:
            x: input tensor (N, C, SpatialDims).
            timesteps: timestep tensor (N,).
            controlnet_cond: controlnet conditioning tensor (N, C, SpatialDims).
            conditioning_scale: conditioning scale.
            context: context tensor (N, 1, ContextDim).
            class_labels: context tensor (N, ).
        """
        # 1. time
        t_emb = get_timestep_embedding(timesteps, self.block_out_channels[0])

        # timesteps does not contain any weights and will always return f32 tensors
        # but time_embedding might actually be running in fp16. so we need to cast here.
        # there might be better ways to encapsulate this.
        t_emb = t_emb.to(dtype=x.dtype)
        emb = self.time_embed(t_emb)

        # 2. class
        if self.num_class_embeds is not None:
            if class_labels is None:
                raise ValueError("class_labels should be provided when num_class_embeds > 0")
            class_emb = self.class_embedding(class_labels)
            class_emb = class_emb.to(dtype=x.dtype)
            emb = emb + class_emb

        # 3. initial convolution
        h = self.conv_in(x)

        # controlnet_cond = self.controlnet_cond_embedding(controlnet_cond)
        controlnet_cond = torch.utils.checkpoint.checkpoint(self.controlnet_cond_embedding,

Review comment (suggestion):
I would suggest splitting the forward method into smaller functional chunks to make it easier to understand and maintain. What do you think?
Something like this: 
```
class CustomControlNet(ControlNet):
    def __init__(self, ...):
        super().__init__(...)

    def forward(self, x, timesteps, controlnet_cond, conditioning_scale=1.0, context=None, class_labels=None):
        emb = self._prepare_time_and_class_embedding(timesteps, class_labels)
        h = self._initial_convolution(x)
        h += self.controlnet_cond_embedding(controlnet_cond)
        h = self._apply_down_blocks(emb, context, h)
        h = self._apply_mid_block(emb, context, h)
        result = self._apply_controlnet_blocks(h, conditioning_scale)
        return result

    def _prepare_time_and_class_embedding(self, timesteps, class_labels):
        ...

    def _initial_convolution(self, x):
        ...

    def _apply_down_blocks(self, emb, context, h):
        ...
        return down_block_res_samples, mid_block_res_sample

    def _apply_mid_block(self, emb, context, h):
        ...

    def _apply_controlnet_blocks(self, h, conditioning_scale):
        ...
        return down_block_res_samples, mid_block_res_sample
```

Code after:
        conditioning_embedding_in_channels: int = 1,
        conditioning_embedding_num_channels: Sequence[int] | None = (16, 32, 96, 256),
        use_checkpointing: bool = True,
    ) -> None:
        super().__init__(
            spatial_dims,
            in_channels,
            num_res_blocks,
            num_channels,
            attention_levels,
            norm_num_groups,
            norm_eps,
            resblock_updown,
            num_head_channels,
            with_conditioning,
            transformer_num_layers,
            cross_attention_dim,
            num_class_embeds,
            upcast_attention,
            use_flash_attention,
            conditioning_embedding_in_channels,
            conditioning_embedding_num_channels,
        )
        self.use_checkpointing = use_checkpointing

    def forward(
        self,
        x: torch.Tensor,
        timesteps: torch.Tensor,
        controlnet_cond: torch.Tensor,
        conditioning_scale: float = 1.0,
        context: torch.Tensor | None = None,
        class_labels: torch.Tensor | None = None,
    ) -> tuple[Sequence[torch.Tensor], torch.Tensor]:
        emb = self._prepare_time_and_class_embedding(x, timesteps, class_labels)
        h = self._apply_initial_convolution(x)
        if self.use_checkpointing:
            controlnet_cond = torch.utils.checkpoint.checkpoint(
                self.controlnet_cond_embedding, controlnet_cond, use_reentrant=False
            )
        else:
            controlnet_cond = self.controlnet_cond_embedding(controlnet_cond)
        h += controlnet_cond
        down_block_res_samples, h = self._apply_down_blocks(emb, context, h)
        h = self._apply_mid_block(emb, context, h)
        down_block_res_samples, mid_block_res_sample = self._apply_controlnet_blocks(h, down_block_res_samples)
        # scaling
        down_block_res_samples = [h * conditioning_scale for h in down_block_res_samples]
        mid_block_res_sample *= conditioning_scale

        return down_block_res_samples, mid_block_res_sample