Update app.py

app.py

@@ -85,6 +85,11 @@ current_model_name = None
 current_pipe = None
 model_lock = threading.Lock()
 
+# Create persistent directory for test images
+PERSISTENT_IMAGE_DIR = "generated_test_images"
+os.makedirs(PERSISTENT_IMAGE_DIR, exist_ok=True)
+print(f"📁 Created persistent image directory: {PERSISTENT_IMAGE_DIR}")
+
 def load_model(model_name="dreamshaper-8"):
     """Thread-safe model loading with HIGH-QUALITY settings"""
     global model_cache, current_model_name, current_pipe
@@ -202,8 +207,35 @@ def generate_high_quality_image(prompt, model_choice="dreamshaper-8", style="chi
         print(f"❌ HQ Generation failed: {str(e)}")
         raise
 
+def save_image_to_persistent_storage(image, prompt, subfolder=""):
+    """Save image to persistent storage that appears in Files tab"""
+    try:
+        # Create filename
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        safe_prompt = "".join(c for c in prompt[:50] if c.isalnum() or c in (' ', '-', '_')).rstrip()
+        filename = f"image_{safe_prompt}_{timestamp}.png"
+        
+        # Create subfolder if specified
+        if subfolder:
+            save_dir = os.path.join(PERSISTENT_IMAGE_DIR, subfolder)
+            os.makedirs(save_dir, exist_ok=True)
+            filepath = os.path.join(save_dir, filename)
+        else:
+            filepath = os.path.join(PERSISTENT_IMAGE_DIR, filename)
+        
+        # Save the image
+        image.save(filepath)
+        print(f"💾 Image saved to persistent storage: {filepath}")
+        
+        # Return relative path for display
+        return filepath
+        
+    except Exception as e:
+        print(f"❌ Failed to save to persistent storage: {e}")
+        return None
+
 def generate_test_image(prompt, model_choice, style_choice):
-    """Generate a single image for testing - FIXED for Hugging Face Spaces"""
+    """Generate a single image for testing and save to persistent storage"""
     try:
         if not prompt.strip():
             return None, "❌ Please enter a prompt"
@@ -221,26 +253,28 @@ def generate_test_image(prompt, model_choice, style_choice):
             negative_prompt
         )
         
-        # Create filename for reference (but don't rely on persistent saving)
-        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-        safe_prompt = "".join(c for c in prompt[:30] if c.isalnum() or c in (' ', '-', '_')).rstrip()
-        filename = f"test_{safe_prompt}_{timestamp}.png"
+        # Save to persistent storage with style-based subfolder
+        saved_path = save_image_to_persistent_storage(image, prompt, style_choice)
         
-        # Try to save to temporary location (optional)
-        try:
-            temp_path = f"/tmp/{filename}"
-            image.save(temp_path)
-            print(f"💾 Temporary save: {temp_path}")
-        except Exception as e:
-            print(f"⚠️ Could not save temporarily: {e}")
+        if saved_path:
+            # Get just the filename for display
+            filename = os.path.basename(saved_path)
+            save_info = f"💾 **Auto-saved to Files tab:** `{PERSISTENT_IMAGE_DIR}/{style_choice}/{filename}`"
+        else:
+            save_info = "⚠️ Could not auto-save to persistent storage"
         
         status_msg = f"""✅ Success! Generated: {prompt}
 
-💾 **Image is ready! Use the DOWNLOAD button below to save it permanently.**
+{save_info}
 
 🎨 Enhanced prompt: {enhanced_prompt}
 
-💡 **Tip:** The download button appears below the image after generation."""
+📁 **Location in Files tab:** 
+- Navigate to: `{PERSISTENT_IMAGE_DIR}/` 
+- Then go to: `{style_choice}/` folder
+- Find your image: `{os.path.basename(saved_path) if saved_path else 'filename'}`
+
+💡 You can also use the download button below the image!"""
         
         return image, status_msg
         
@@ -341,11 +375,38 @@ async def get_job_status_endpoint(job_id: str):
         updated_at=job_data["updated_at"]
     )
 
+@app.get("/api/list-test-images")
+async def list_test_images():
+    """API endpoint to list all generated test images"""
+    try:
+        images_list = []
+        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)
+                    rel_path = os.path.relpath(full_path, PERSISTENT_IMAGE_DIR)
+                    images_list.append({
+                        "filename": file,
+                        "path": rel_path,
+                        "full_path": full_path,
+                        "size": os.path.getsize(full_path) if os.path.exists(full_path) else 0,
+                        "created": os.path.getctime(full_path) if os.path.exists(full_path) else 0
+                    })
+        
+        return {
+            "status": "success",
+            "image_count": len(images_list),
+            "image_directory": PERSISTENT_IMAGE_DIR,
+            "images": sorted(images_list, key=lambda x: x["created"], reverse=True)
+        }
+    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 instantly!**")
+        gr.Markdown("**Generate studio-quality images that auto-save to your Files tab!**")
         
         with gr.Row():
             with gr.Column(scale=1):
@@ -369,11 +430,12 @@ def create_test_interface():
                 
                 generate_btn = gr.Button("✨ Generate Image", variant="primary", size="lg")
                 
-                gr.Markdown("### 💡 Tips")
-                gr.Markdown("""
-                - Be descriptive: **"A cute baby dragon with green scales reading under a magical tree"**
-                - Download images using the button below each generated image
-                - Test different models and styles for varied results
+                gr.Markdown("### 💡 Auto-Save Feature")
+                gr.Markdown(f"""
+                **All images are automatically saved to:**
+                - `{PERSISTENT_IMAGE_DIR}/[style]/` in your Files tab
+                - Organized by art style for easy browsing
+                - Also available via download button below
                 """)
             
             with gr.Column(scale=2):
@@ -384,8 +446,8 @@ def create_test_interface():
                     show_share_button=True
                 )
                 status_output = gr.Textbox(
-                    label="Status", 
-                    lines=4,
+                    label="Status & File Location", 
+                    lines=5,
                     show_copy_button=True
                 )
         
@@ -404,10 +466,42 @@ def create_test_interface():
                 outputs=prompt_input
             )
         
+        # Image browser section
+        with gr.Accordion("📁 Browse Generated Images", open=False):
+            gr.Markdown(f"### Images saved in `{PERSISTENT_IMAGE_DIR}/`")
+            refresh_btn = gr.Button("🔄 Refresh File List")
+            file_list = gr.File(
+                label="Generated Images",
+                file_count="multiple",
+                type="filepath",
+                height=300
+            )
+            
+            def refresh_file_list():
+                """Refresh the list of generated images"""
+                try:
+                    image_files = []
+                    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
+                except Exception as e:
+                    print(f"Error refreshing file list: {e}")
+                    return []
+            
+            refresh_btn.click(fn=refresh_file_list, outputs=file_list)
+            # Auto-refresh on load
+            demo.load(fn=refresh_file_list, outputs=file_list)
+        
         generate_btn.click(
             fn=generate_test_image,
             inputs=[prompt_input, model_dropdown, style_dropdown],
             outputs=[image_output, status_output]
+        ).then(
+            fn=refresh_file_list,
+            outputs=file_list
         )
     
     return demo
@@ -419,10 +513,32 @@ gradio_app = create_test_interface()
 async def root():
     return {"message": "Storybook Generator API + UI is running!"}
 
+@app.get("/files")
+async def files_info():
+    """Endpoint to check file structure"""
+    try:
+        file_structure = {}
+        for root, dirs, files in os.walk(PERSISTENT_IMAGE_DIR):
+            rel_root = os.path.relpath(root, PERSISTENT_IMAGE_DIR)
+            file_structure[rel_root] = {
+                "files": files,
+                "file_count": len(files)
+            }
+        
+        return {
+            "persistent_directory": PERSISTENT_IMAGE_DIR,
+            "exists": os.path.exists(PERSISTENT_IMAGE_DIR),
+            "file_structure": file_structure
+        }
+    except Exception as e:
+        return {"error": str(e)}
+
 # For Hugging Face Spaces
 def get_app():
     return app
 
 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")
     gradio_app.launch(server_name="0.0.0.0", server_port=7860)