control working

Author: miguelmartin75

src/diffusers/models/controlnets/controlnet_cosmos.py

@@ -135,7 +135,7 @@ def forward(
         scales = self._expand_conditioning_scale(conditioning_scale)
         result = []
         for block_idx, (block, scale) in enumerate(zip(self.control_blocks, scales)):
-            control_hidden_states = block(
+            control_hidden_states, control_proj = block(
                 hidden_states=control_hidden_states,
                 encoder_hidden_states=encoder_hidden_states,
                 embedded_timestep=embedded_timestep,
@@ -147,5 +147,5 @@ def forward(
                 block_idx=block_idx,
                 latents=latents,
             )
-            result.append(control_hidden_states * scale)
+            result.append(control_proj * scale)
         return result

src/diffusers/models/transformers/transformer_cosmos.py

@@ -419,7 +419,7 @@ def forward(
         controlnet_residual: Optional[torch.Tensor] = None,
         latents: Optional[torch.Tensor] = None,
         block_idx: Optional[int] = None,
-    ) -> torch.Tensor:
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
         if self.before_proj is not None:
             hidden_states = self.before_proj(hidden_states) + latents
             print(f"before_proj, block_idx={block_idx}")
@@ -445,8 +445,10 @@ def forward(
         hidden_states = hidden_states + gate * ff_output
 
         if self.after_proj is not None:
-            hidden_states = self.after_proj(hidden_states)
+            assert controlnet_residual is None
+            hs_proj = self.after_proj(hidden_states)
             print(f"after_proj, block_idx={block_idx}")
+            return hidden_states, hs_proj
 
         if controlnet_residual is not None:
             # NOTE: this is assumed to be scaled by the controlnet
@@ -846,7 +848,6 @@ def _forward(self, prepared_inputs: Dict[str, Any], block_controlnet_hidden_stat
                     prepared_inputs["extra_pos_emb"],
                     prepared_inputs["attention_mask"],
                     controlnet_residual,
-                    latents,
                 )
             else:
                 hidden_states = block(
@@ -858,7 +859,6 @@ def _forward(self, prepared_inputs: Dict[str, Any], block_controlnet_hidden_stat
                     prepared_inputs["extra_pos_emb"],
                     prepared_inputs["attention_mask"],
                     controlnet_residual,
-                    latents,
                 )
 
         temb = prepared_inputs["temb"]

src/diffusers/pipelines/cosmos/pipeline_cosmos2_5_transfer.py

@@ -76,15 +76,15 @@ def retrieve_latents(
 
 # TODO: move this to a utility module aka Transfer2_5 model ?
 def transfer2_5_forward(
-    transformer,
-    controlnet,
-    in_latents,
-    controls_latents,
-    controls_conditioning_scale,
-    in_timestep,
-    encoder_hidden_states,
-    cond_mask,
-    padding_mask,
+    transformer: CosmosTransformer3DModel,
+    controlnet: CosmosControlNetModel,
+    in_latents: torch.Tensor,
+    controls_latents: torch.Tensor,
+    controls_conditioning_scale: list[float],
+    in_timestep: torch.Tensor,
+    encoder_hidden_states: tuple[torch.Tensor | None, torch.Tensor | None] | None,
+    cond_mask: torch.Tensor,
+    padding_mask: torch.Tensor,
 ):
     control_blocks = None
     prepared_inputs = transformer.prepare_inputs(
@@ -97,7 +97,7 @@ def transfer2_5_forward(
     if controls_latents is not None:
         control_blocks = controlnet(
             controls_latents=controls_latents,
-            latents=in_latents,
+            latents=prepared_inputs["hidden_states"],
             conditioning_scale=controls_conditioning_scale,
             condition_mask=cond_mask,
             padding_mask=padding_mask,