align latents to transformer

app.py

@@ -46,10 +46,24 @@ def infer(prompt, init_image=None, seed=42, randomize_seed=False, width=1024, he
         init_image = init_image.convert("RGB")
         init_image = preprocess_image(init_image, vae_image_size)
         latents = encode_image(init_image, pipe.vae)
+
+        # Debug: Print the shape of the latents after encoding
+        print(f"Latents shape after encoding: {latents.shape}")
+
         # Ensure latents are correctly shaped and adjusted
         latents = torch.nn.functional.interpolate(latents, size=(height // 8, width // 8))
-        latents = latents.view(1, -1, height // 8, width // 8)
         
+        # Convert latent channels to 64 as expected by the transformer
+        latent_channels = pipe.vae.config.latent_channels
+        if latent_channels != 64:
+            latents = torch.nn.Conv2d(latent_channels, 64, kernel_size=1).to(device)(latents)
+
+        # Reshape latents to match the transformer's input expectations
+        latents = latents.view(1, 64, height // 8, width // 8)
+
+        # Debug: Print the shape of the latents after reshaping
+        print(f"Latents shape after reshaping: {latents.shape}")
+
         image = pipe(
             prompt=prompt,
             height=height,
@@ -73,7 +87,6 @@ def infer(prompt, init_image=None, seed=42, randomize_seed=False, width=1024, he
     return image, seed
 
 
-
 # Define example prompts
 examples = [
     "a tiny astronaut hatching from an egg on the moon",