Upload 6 files

app.py

@@ -552,7 +552,27 @@ def main():
                 image_cols = st.columns(min(len(images), 3))
                 for i, img_path in enumerate(images):
                     with image_cols[i % len(image_cols)]:
-                        st.image(img_path, caption=f"Image {i+1}", use_column_width=True)
+                        try:
+                            # Verify image exists and is valid
+                            if os.path.exists(img_path):
+                                # Try to open and verify the image
+                                from PIL import Image
+                                try:
+                                    img = Image.open(img_path)
+                                    # Convert to RGB if needed
+                                    if img.mode != "RGB":
+                                        img = img.convert("RGB")
+                                    # Save as JPEG to ensure compatibility
+                                    safe_path = f"temp/safe_image_{int(time.time() * 1000)}_{i}.jpg"
+                                    img.save(safe_path, format="JPEG", quality=95)
+                                    # Display the safe image
+                                    st.image(safe_path, caption=f"Image {i+1}", use_container_width=True)
+                                except Exception as e:
+                                    st.error(f"Error loading image {i+1}: {str(e)}")
+                            else:
+                                st.warning(f"Image {i+1} not found")
+                        except Exception as e:
+                            st.error(f"Error displaying image {i+1}: {str(e)}")
                 
                 # Force garbage collection after image generation
                 if memory_optimization or apply_vram_optimization:

image_generator.py

@@ -188,9 +188,20 @@ class ImageGenerator:
                         height=min(self.target_size[1], 512)  # Limit size for Hugging Face
                     ).images[0]
                 
-                # Save the image to a temporary file
-                output_path = f"temp/image_{int(time.time() * 1000)}.png"
-                image.save(output_path)
+                # Save the image to a temporary file with explicit format
+                output_path = f"temp/image_{int(time.time() * 1000)}.jpg"
+                image = image.convert("RGB")  # Ensure image is in RGB mode
+                image.save(output_path, format="JPEG", quality=95)  # Use JPEG format explicitly
+                
+                # Verify the image was saved correctly
+                try:
+                    from PIL import Image
+                    test_load = Image.open(output_path)
+                    test_load.verify()  # Verify image is valid
+                    test_load.close()
+                except Exception as e:
+                    st.error(f"Image verification failed: {str(e)}. Using fallback.")
+                    return self.create_fallback_image(prompt)
                 
                 # Force garbage collection after inference
                 gc.collect()
@@ -248,9 +259,20 @@ class ImageGenerator:
                 height=512
             ).images[0]
             
-            # Save the image
-            output_path = f"temp/retry_image_{int(time.time() * 1000)}.png"
-            image.save(output_path)
+            # Save the image with explicit format
+            output_path = f"temp/retry_image_{int(time.time() * 1000)}.jpg"
+            image = image.convert("RGB")  # Ensure image is in RGB mode
+            image.save(output_path, format="JPEG", quality=95)  # Use JPEG format explicitly
+            
+            # Verify the image was saved correctly
+            try:
+                from PIL import Image
+                test_load = Image.open(output_path)
+                test_load.verify()  # Verify image is valid
+                test_load.close()
+            except Exception as e:
+                st.error(f"Image verification failed: {str(e)}. Using fallback.")
+                return self.create_fallback_image(prompt)
             
             return output_path
         except Exception as e: