Update pipeline.py

pipeline.py

@@ -240,7 +240,7 @@ class Flex2Pipeline(FluxControlPipeline):
         if control_image is None:
             control_latents = torch.zeros(
                 batch_size * num_images_per_prompt,
-                3,
+                16,
                 latent_height,
                 latent_width,
                 device=device,
@@ -261,12 +261,11 @@ class Flex2Pipeline(FluxControlPipeline):
         
         # apply control strength
         control_latents = control_latents * control_strength
-        print("control_latents", control_latents.shape)
         
         if inpaint_image is None and inpaint_mask is None:
             inpaint_latents = torch.zeros(
                 batch_size * num_images_per_prompt,
-                3,
+                16,
                 latent_height,
                 latent_width,
                 device=device,
@@ -282,7 +281,7 @@ class Flex2Pipeline(FluxControlPipeline):
             )
         else:
             print("inpaint_image.shape",inpaint_image.size)
-            print("inpaint_mask.shape",inpaint_mask.size)
+            print("inpaint_mask.shape",inpaint_mask.shape)
             inpaint_image = self.prepare_image(
                 image=inpaint_image,
                 width=width,
@@ -294,7 +293,6 @@ class Flex2Pipeline(FluxControlPipeline):
             )
             inpaint_image = self.vae.encode(inpaint_image).latent_dist.sample(generator=generator)
             inpaint_latents = (inpaint_image - self.vae.config.shift_factor) * self.vae.config.scaling_factor
-            print("inpaint_latents", inpaint_latents.shape)
             height_inpaint_image, width_inpaint_image = inpaint_image.shape[2:]
             
             inpaint_mask = self.prepare_image(
@@ -310,7 +308,7 @@ class Flex2Pipeline(FluxControlPipeline):
             inpaint_mask = inpaint_mask[:, 0:1, :, :] * 0.5 + 0.5
             # resize to match height_inpaint_image and width_inpaint_image
             inpaint_latents_mask = F.interpolate(inpaint_mask, size=(height_inpaint_image, width_inpaint_image), mode="bilinear", align_corners=False)
-            print("inpaint_latents_mask", inpaint_latents_mask.shape)
+        
         # apply inverted mask to inpaint latents
         inpaint_latents = inpaint_latents * (1 - inpaint_latents_mask)
         
@@ -443,4 +441,6 @@ class Flex2Pipeline(FluxControlPipeline):
         if not return_dict:
             return (image,)
 
-        return FluxPipelineOutput(images=image)
+        return FluxPipelineOutput(images=image)
+
+