Update app.py

app.py

@@ -24,9 +24,14 @@ if st.button("Generate Image") and text:
     # Process image and get CLIP features for image
     image_features = clip_processor(images=example_image, return_tensors="pt", padding=True)
 
+    # Ensure the dimensions of pixel_values are the same for text and image features
+    max_len = max(text_features.pixel_values.shape[1], image_features.pixel_values.shape[1])
+    text_features.pixel_values = torch.nn.functional.pad(text_features.pixel_values, (0, max_len - text_features.pixel_values.shape[1]))
+    image_features.pixel_values = torch.nn.functional.pad(image_features.pixel_values, (0, max_len - image_features.pixel_values.shape[1]))
+
     # Concatenate text and image features
     combined_features = {
-        "pixel_values": torch.cat([text_features.pixel_values, image_features.pixel_values], dim=-1)
+        "pixel_values": torch.cat([text_features.pixel_values, image_features.pixel_values], dim=1)
     }
 
     # Forward pass through CLIP
@@ -38,6 +43,3 @@ if st.button("Generate Image") and text:
 
     # Display the generated image
     st.image(generated_image, caption="Generated Image", use_column_width=True)
-
-    # Save the generated image to a file (change the file path as needed)
-    generated_image.save("generated_image.jpg")