Update app.py

app.py

@@ -68,22 +68,17 @@ if st.button("Generate"):
     with st.spinner("Generating images..."):        
         images = pipe(prompt, guidance_scale=guidance_scale, num_inference_steps=64).images        
 
-def process_image_data(image_data):
-    # Assume image_data has a nested structure or is a custom object
-    # Extract the raw pixel data and reshape as needed
-    # Example: 
-    if isinstance(image_data, torch.Tensor):  
-        return image_data.numpy()  
-    # ... handle other potential data structures
-
-# ... (Inside your generation loop)
-pil_images = []
-for image in images:
-    processed_data = process_image_data(image)
-    pil_image = Image.fromarray(processed_data)
-    resized_image = pil_image.resize((256, 256)) 
-    pil_images.append(resized_image)
-     
-gif_path = export_to_gif(pil_images, "shark_3d.gif") 
-st.image(pil_images[0])
-st.success("GIF saved as shark_3d.gif")
+        # Process images for PIL conversion (This will be customized)
+        pil_images = []
+        for image in images:
+            processed_image = process_image_for_pil(image) 
+            pil_images.append(processed_image)
+
+        # Resize Images (Optional)
+        if should_resize():  # Add a function to control if resizing is needed 
+            for i in range(len(pil_images)):
+                pil_images[i] = pil_images[i].resize((256, 256)) 
+
+        gif_path = export_to_gif(pil_images, "shark_3d.gif") 
+        st.image(pil_images[0])  
+        st.success("GIF saved as shark_3d.gif")