added more console logging

app.py

@@ -29,32 +29,13 @@ def preprocess_image(image, image_size):
     return image
 
 def check_shapes(latents):
-    # Get the shape of the latents
-    latent_shape = latents.shape
-    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])
-    print(f"Expected transformer input shape: {expected_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]}")
-    except AttributeError as e:
-        print(f"Unable to access transformer weights: {e}")
+    print(f"Latent shape: {latents.shape}")
+    if len(latents.shape) == 4:
+        print(f"Expected transformer input shape: {(1, latents.shape[1] * latents.shape[2] * latents.shape[3])}")
+    elif len(latents.shape) == 2:
+        print(f"Reshaped latent shape: {latents.shape}")
+    else:
+        print(f"Unexpected latent shape: {latents.shape}")
 
 @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)):
@@ -93,6 +74,15 @@ def infer(prompt, init_image=None, seed=42, randomize_seed=False, width=1024, he
             # Check shapes after reshaping
             check_shapes(latents)
 
+            # Print the type and shape of each argument
+            print(f"prompt type: {type(prompt)}, value: {prompt}")
+            print(f"height type: {type(height)}, value: {height}")
+            print(f"width type: {type(width)}, value: {width}")
+            print(f"num_inference_steps type: {type(num_inference_steps)}, value: {num_inference_steps}")
+            print(f"generator type: {type(generator)}")
+            print(f"guidance_scale type: {type(0.0)}, value: 0.0")
+            print(f"latents type: {type(latents)}, shape: {latents.shape}")
+
             image = pipe(
                 prompt=prompt,
                 height=height,
@@ -106,6 +96,8 @@ def infer(prompt, init_image=None, seed=42, randomize_seed=False, width=1024, he
         return image, seed
     except Exception as e:
         print(f"Error during inference: {e}")
+        import traceback
+        traceback.print_exc()
         return Image.new("RGB", (width, height), (255, 0, 0)), seed  # Red fallback image
 
 # Gradio interface setup