Update app.py

app.py

@@ -234,11 +234,44 @@ def save_image_to_persistent_storage(image, prompt, subfolder=""):
         print(f"❌ Failed to save to persistent storage: {e}")
         return None
 
+def delete_image_file(filepath):
+    """Delete an image file from persistent storage"""
+    try:
+        if os.path.exists(filepath):
+            os.remove(filepath)
+            print(f"🗑️ Deleted image: {filepath}")
+            return True, f"✅ Successfully deleted: {os.path.basename(filepath)}"
+        else:
+            return False, f"❌ File not found: {filepath}"
+    except Exception as e:
+        return False, f"❌ Error deleting file: {str(e)}"
+
+def get_storage_info():
+    """Get information about storage usage"""
+    try:
+        total_size = 0
+        file_count = 0
+        
+        for root, dirs, files in os.walk(PERSISTENT_IMAGE_DIR):
+            for file in files:
+                filepath = os.path.join(root, file)
+                if os.path.exists(filepath):
+                    total_size += os.path.getsize(filepath)
+                    file_count += 1
+        
+        return {
+            "total_files": file_count,
+            "total_size_mb": round(total_size / (1024 * 1024), 2),
+            "directory": PERSISTENT_IMAGE_DIR
+        }
+    except Exception as e:
+        return {"error": str(e)}
+
 def generate_test_image(prompt, model_choice, style_choice):
     """Generate a single image for testing and save to persistent storage"""
     try:
         if not prompt.strip():
-            return None, "❌ Please enter a prompt"
+            return None, "❌ Please enter a prompt", None
         
         print(f"🎨 Generating test image with prompt: {prompt}")
         
@@ -276,12 +309,12 @@ def generate_test_image(prompt, model_choice, style_choice):
 
 💡 You can also use the download button below the image!"""
         
-        return image, status_msg
+        return image, status_msg, saved_path
         
     except Exception as e:
         error_msg = f"❌ Generation failed: {str(e)}"
         print(error_msg)
-        return None, error_msg
+        return None, error_msg, None
 
 # JOB MANAGEMENT FUNCTIONS
 def create_job(story_request: StorybookRequest) -> str:
@@ -393,21 +426,77 @@ async def list_test_images():
                         "created": os.path.getctime(full_path) if os.path.exists(full_path) else 0
                     })
         
+        storage_info = get_storage_info()
+        
         return {
             "status": "success",
             "image_count": len(images_list),
             "image_directory": PERSISTENT_IMAGE_DIR,
+            "storage_info": storage_info,
             "images": sorted(images_list, key=lambda x: x["created"], reverse=True)
         }
     except Exception as e:
         return {"status": "error", "message": str(e)}
 
+@app.delete("/api/delete-image/{image_path:path}")
+async def delete_image(image_path: str):
+    """API endpoint to delete an image"""
+    try:
+        # Security check: ensure the path is within our image directory
+        full_path = os.path.join(PERSISTENT_IMAGE_DIR, image_path)
+        if not full_path.startswith(os.path.abspath(PERSISTENT_IMAGE_DIR)):
+            return {"status": "error", "message": "Invalid path"}
+        
+        success, message = delete_image_file(full_path)
+        return {"status": "success" if success else "error", "message": message}
+    except Exception as e:
+        return {"status": "error", "message": str(e)}
+
+@app.delete("/api/clear-all-images")
+async def clear_all_images():
+    """API endpoint to delete all generated images"""
+    try:
+        deleted_count = 0
+        error_count = 0
+        
+        for root, dirs, files in os.walk(PERSISTENT_IMAGE_DIR):
+            for file in files:
+                if file.endswith(('.png', '.jpg', '.jpeg')):
+                    filepath = os.path.join(root, file)
+                    success, _ = delete_image_file(filepath)
+                    if success:
+                        deleted_count += 1
+                    else:
+                        error_count += 1
+        
+        return {
+            "status": "success",
+            "message": f"Deleted {deleted_count} images, {error_count} errors",
+            "deleted_count": deleted_count,
+            "error_count": error_count
+        }
+    except Exception as e:
+        return {"status": "error", "message": str(e)}
+
 # GRADIO INTERFACE
 def create_test_interface():
     with gr.Blocks(title="High-Quality Image Generator", theme="soft") as demo:
         gr.Markdown("# 🎨 High-Quality Image Generator")
         gr.Markdown("**Generate studio-quality images that auto-save to your Files tab!**")
         
+        # Storage info display
+        storage_info = gr.Textbox(
+            label="📊 Storage Information",
+            interactive=False,
+            lines=2
+        )
+        
+        def update_storage_info():
+            info = get_storage_info()
+            if "error" not in info:
+                return f"📁 Storage: {info['total_files']} images, {info['total_size_mb']} MB used"
+            return "📁 Storage: Unable to calculate"
+        
         with gr.Row():
             with gr.Column(scale=1):
                 model_dropdown = gr.Dropdown(
@@ -441,13 +530,19 @@ def create_test_interface():
             with gr.Column(scale=2):
                 image_output = gr.Image(
                     label="Generated Image", 
-                    height=500, 
+                    height=400, 
                     show_download_button=True,
                     show_share_button=True
                 )
+                
+                # Delete button for current image
+                current_file_path = gr.State()  # Hidden state to store file path
+                delete_btn = gr.Button("🗑️ Delete This Image", variant="stop", size="sm")
+                delete_status = gr.Textbox(label="Delete Status", interactive=False, lines=2)
+                
                 status_output = gr.Textbox(
                     label="Status & File Location", 
-                    lines=5,
+                    lines=4,
                     show_copy_button=True
                 )
         
@@ -466,43 +561,129 @@ def create_test_interface():
                 outputs=prompt_input
             )
         
-        # Image browser section
-        with gr.Accordion("📁 Browse Generated Images", open=False):
+        # Image browser section with delete functionality
+        with gr.Accordion("📁 Manage Generated Images", open=False):
             gr.Markdown(f"### Images saved in `{PERSISTENT_IMAGE_DIR}/`")
-            refresh_btn = gr.Button("🔄 Refresh File List")
-            file_list = gr.File(
+            
+            with gr.Row():
+                refresh_btn = gr.Button("🔄 Refresh List")
+                clear_all_btn = gr.Button("🗑️ Clear All Images", variant="stop")
+            
+            file_list = gr.Gallery(
                 label="Generated Images",
-                file_count="multiple",
-                type="filepath",
-                height=300
+                show_label=True,
+                elem_id="gallery",
+                columns=3,
+                height="auto"
             )
             
+            selected_file_info = gr.Textbox(label="Selected Image Info", interactive=False)
+            delete_selected_btn = gr.Button("🗑️ Delete Selected Image", variant="stop")
+            delete_selected_status = gr.Textbox(label="Delete Status", interactive=False)
+            
             def refresh_file_list():
                 """Refresh the list of generated images"""
                 try:
                     image_files = []
+                    image_info = []
+                    
                     for root, dirs, files in os.walk(PERSISTENT_IMAGE_DIR):
                         for file in files:
                             if file.endswith(('.png', '.jpg', '.jpeg')):
                                 full_path = os.path.join(root, file)
-                                image_files.append(full_path)
-                    return image_files
+                                if os.path.exists(full_path):
+                                    image_files.append(full_path)
+                                    image_info.append(f"📄 {file}\n📁 {os.path.relpath(root, PERSISTENT_IMAGE_DIR)}")
+                    
+                    return image_files, update_storage_info()
                 except Exception as e:
                     print(f"Error refreshing file list: {e}")
-                    return []
+                    return [], "❌ Error loading images"
             
-            refresh_btn.click(fn=refresh_file_list, outputs=file_list)
-            # Auto-refresh on load
-            demo.load(fn=refresh_file_list, outputs=file_list)
+            def get_file_info(evt: gr.SelectData):
+                """Get info about selected file"""
+                if evt.value:
+                    filepath = evt.value
+                    filename = os.path.basename(filepath)
+                    size_kb = os.path.getsize(filepath) / 1024 if os.path.exists(filepath) else 0
+                    return f"Selected: {filename}\nSize: {size_kb:.1f} KB\nPath: {filepath}"
+                return "No image selected"
+            
+            def delete_selected_file(selected_files):
+                """Delete the selected file"""
+                if not selected_files:
+                    return "❌ No image selected", None, update_storage_info()
+                
+                filepath = selected_files[0] if isinstance(selected_files, list) else selected_files
+                success, message = delete_image_file(filepath)
+                
+                # Refresh the file list
+                new_files, storage_info = refresh_file_list()
+                return message, new_files, storage_info
+            
+            def clear_all_images_func():
+                """Clear all images"""
+                try:
+                    deleted_count = 0
+                    for root, dirs, files in os.walk(PERSISTENT_IMAGE_DIR):
+                        for file in files:
+                            if file.endswith(('.png', '.jpg', '.jpeg')):
+                                filepath = os.path.join(root, file)
+                                success, _ = delete_image_file(filepath)
+                                if success:
+                                    deleted_count += 1
+                    
+                    new_files, storage_info = refresh_file_list()
+                    return f"✅ Deleted {deleted_count} images", new_files, storage_info
+                except Exception as e:
+                    return f"❌ Error: {str(e)}", None, update_storage_info()
+            
+            # Connect events
+            refresh_btn.click(fn=refresh_file_list, outputs=[file_list, storage_info])
+            clear_all_btn.click(fn=clear_all_images_func, outputs=[delete_selected_status, file_list, storage_info])
+            file_list.select(fn=get_file_info, outputs=selected_file_info)
+            delete_selected_btn.click(
+                fn=delete_selected_file, 
+                inputs=file_list,
+                outputs=[delete_selected_status, file_list, storage_info]
+            )
+        
+        # Delete current image functionality
+        def delete_current_image(filepath):
+            """Delete the currently displayed image"""
+            if not filepath:
+                return "❌ No image to delete", None, None, update_storage_info()
+            
+            success, message = delete_image_file(filepath)
+            new_files, storage_info = refresh_file_list()
+            
+            if success:
+                return message, None, "✅ Image deleted successfully!", storage_info
+            else:
+                return message, None, "❌ Delete failed", storage_info
         
+        # Connect main generate button
         generate_btn.click(
             fn=generate_test_image,
             inputs=[prompt_input, model_dropdown, style_dropdown],
-            outputs=[image_output, status_output]
+            outputs=[image_output, status_output, current_file_path]
         ).then(
             fn=refresh_file_list,
-            outputs=file_list
+            outputs=[file_list, storage_info]
         )
+        
+        # Connect delete button
+        delete_btn.click(
+            fn=delete_current_image,
+            inputs=current_file_path,
+            outputs=[delete_status, image_output, status_output, storage_info]
+        ).then(
+            fn=refresh_file_list,
+            outputs=[file_list, storage_info]
+        )
+        
+        # Initialize on load
+        demo.load(fn=refresh_file_list, outputs=[file_list, storage_info])
     
     return demo
 
@@ -528,7 +709,8 @@ async def files_info():
         return {
             "persistent_directory": PERSISTENT_IMAGE_DIR,
             "exists": os.path.exists(PERSISTENT_IMAGE_DIR),
-            "file_structure": file_structure
+            "file_structure": file_structure,
+            "storage_info": get_storage_info()
         }
     except Exception as e:
         return {"error": str(e)}
@@ -541,4 +723,5 @@ if __name__ == "__main__":
     print("🚀 Starting High-Quality Storybook Generator...")
     print(f"📁 Persistent image directory: {PERSISTENT_IMAGE_DIR}")
     print("✅ Images will auto-save to your Files tab")
+    print("🗑️  Delete functionality enabled for storage management")
     gradio_app.launch(server_name="0.0.0.0", server_port=7860)