1

Author: zRzRzRzRzRzRzR

examples/cogview4-control/train_control_cogview4.py

@@ -804,6 +804,7 @@ def main(args):
         cogview4_transformer.patch_embed.proj = new_linear
 
     assert torch.all(cogview4_transformer.patch_embed.proj.weight[:, initial_input_channels:].data == 0)
+    cogview4_transformer.register_to_config(in_channels=cogview4_transformer.config.in_channels * 2, out_channels=initial_input_channels)
 
     if args.only_target_transformer_blocks:
         cogview4_transformer.patch_embed.proj.requires_grad_(True)

src/diffusers/pipelines/cogview4/pipeline_cogview4_control.py

@@ -579,7 +579,7 @@ def __call__(
         )
 
         # Prepare latents
-        latent_channels = self.transformer.config.in_channels
+        latent_channels = self.transformer.config.in_channels // 2
 
         control_image = self.prepare_image(
             image=control_image,