remove latent flattening

app.py

@@ -18,7 +18,6 @@ pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", tor
 
 def preprocess_image(image, image_size):
     print(f"Preprocessing image to size: {image_size}x{image_size}")
-    # Preprocess the image for the VAE
     preprocess = transforms.Compose([
         transforms.Resize((image_size, image_size)),  # Use model-specific size
         transforms.ToTensor(),
@@ -30,7 +29,6 @@ def preprocess_image(image, image_size):
 
 def encode_image(image, vae):
     print("Encoding image using the VAE")
-    # Encode the image using the VAE
     with torch.no_grad():
         latents = vae.encode(image).latent_dist.sample() * 0.18215
     print(f"Latents shape after encoding: {latents.shape}")
@@ -72,9 +70,16 @@ def infer(prompt, init_image=None, seed=42, randomize_seed=False, width=1024, he
         latents = latents.view(1, 64, height // 8, width // 8)
         print(f"Latents shape after reshaping: {latents.shape}")
 
-        # Flatten the latents if required by the transformer
-        latents = latents.flatten(start_dim=1)
-        print(f"Latents shape after flattening: {latents.shape}")
+        # Avoid flattening, ensure latents are in the expected shape for the transformer
+        # Adding extra debug to understand what transformer expects
+        try:
+            print("Calling the transformer with latents")
+            # Dummy call to transformer to understand the shape requirement
+            _ = pipe.transformer(latents)
+            print("Transformer call succeeded")
+        except Exception as e:
+            print(f"Transformer call failed with error: {e}")
+            raise
 
         print("Calling the diffusion pipeline with latents")
         image = pipe(
@@ -103,6 +108,7 @@ def infer(prompt, init_image=None, seed=42, randomize_seed=False, width=1024, he
 
 
 
+
 # Define example prompts
 examples = [
     "a tiny astronaut hatching from an egg on the moon",