trying to match the latent shapes

app.py

@@ -34,22 +34,27 @@ def check_shapes(latents):
     print(f"Latent shape: {latent_shape}")
 
     # Get the expected shape for the transformer input
-    expected_shape = (1, latent_shape[1] * latent_shape[2], latent_shape[3])
+    expected_shape = (1, latent_shape[1] * latent_shape[2] * latent_shape[3])
     print(f"Expected transformer input shape: {expected_shape}")
 
-    # Get the shape of the transformer's weight matrix
-    if hasattr(pipe.transformer, 'text_model'):
-        weight_shape = pipe.transformer.text_model.encoder.layers[0].self_attn.q_proj.weight.shape
-    else:
-        weight_shape = pipe.transformer.encoder.layers[0].self_attn.q_proj.weight.shape
-    print(f"Transformer weight shape: {weight_shape}")
+    # Try to get the shape of the transformer's weight matrix
+    try:
+        # Assuming the first layer of the transformer has a linear projection
+        if hasattr(pipe.transformer, 'blocks'):
+            weight_shape = pipe.transformer.blocks[0].attn.to_q.weight.shape
+        else:
+            print("Unable to determine transformer weight shape.")
+            return
+        print(f"Transformer weight shape: {weight_shape}")
 
-    # Check if the shapes are compatible for matrix multiplication
-    if expected_shape[1] == weight_shape[1]:
-        print("Shapes are compatible for matrix multiplication.")
-    else:
-        print("Warning: Shapes are not compatible for matrix multiplication.")
-        print(f"Expected: {expected_shape[1]}, Got: {weight_shape[1]}")
+        # Check if the shapes are compatible for matrix multiplication
+        if expected_shape[1] == weight_shape[1]:
+            print("Shapes are compatible for matrix multiplication.")
+        else:
+            print("Warning: Shapes are not compatible for matrix multiplication.")
+            print(f"Expected: {expected_shape[1]}, Got: {weight_shape[1]}")
+    except AttributeError as e:
+        print(f"Unable to access transformer weights: {e}")
 
 @spaces.GPU()
 def infer(prompt, init_image=None, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=4, progress=gr.Progress(track_tqdm=True)):
@@ -83,7 +88,7 @@ def infer(prompt, init_image=None, seed=42, randomize_seed=False, width=1024, he
             check_shapes(latents)
             
             # Reshape latents to match the expected input shape of the transformer
-            latents = latents.permute(0, 2, 3, 1).contiguous().view(1, -1, pipe.vae.config.latent_channels)
+            latents = latents.reshape(1, -1)
 
             # Check shapes after reshaping
             check_shapes(latents)
@@ -103,7 +108,6 @@ def infer(prompt, init_image=None, seed=42, randomize_seed=False, width=1024, he
         print(f"Error during inference: {e}")
         return Image.new("RGB", (width, height), (255, 0, 0)), seed  # Red fallback image
 
-
 # Gradio interface setup
 with gr.Blocks() as demo:
     with gr.Row():