Update pipeline.py

pipeline.py

@@ -261,6 +261,7 @@ 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(
@@ -293,6 +294,7 @@ 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(
@@ -308,7 +310,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)