Update app.py

app.py

@@ -23,6 +23,9 @@ import random
 # External OCI API URL
 OCI_API_BASE_URL = "https://yukee1992-oci-story-book.hf.space"
 
+# Create images directory if it doesn't exist
+os.makedirs("generated_test_images", exist_ok=True)
+
 # Initialize FastAPI app
 app = FastAPI(title="Storybook Generator API")
 
@@ -81,7 +84,7 @@ class JobStatusResponse(BaseModel):
     created_at: float
     updated_at: float
 
-# HIGH-QUALITY MODEL SELECTION (from your high-quality app.py)
+# HIGH-QUALITY MODEL SELECTION
 MODEL_CHOICES = {
     "dreamshaper-8": "lykon/dreamshaper-8",  # Excellent for fantasy/artistic
     "realistic-vision": "SG161222/Realistic_Vision_V5.1",  # Photorealistic
@@ -127,7 +130,7 @@ def load_model(model_name="dreamshaper-8"):
         try:
             model_id = MODEL_CHOICES.get(model_name, "lykon/dreamshaper-8")
             
-            # HIGH-QUALITY LOADING (from your high-quality app.py)
+            # HIGH-QUALITY LOADING
             pipe = StableDiffusionPipeline.from_pretrained(
                 model_id, 
                 torch_dtype=torch.float32,
@@ -161,7 +164,7 @@ print("🚀 Initializing HIGH-QUALITY Storybook Generator...")
 load_model("dreamshaper-8")
 print("✅ HIGH-QUALITY Model loaded and ready!")
 
-# PROFESSIONAL PROMPT ENGINEERING (from your high-quality app.py)
+# PROFESSIONAL PROMPT ENGINEERING
 def enhance_prompt(prompt, style="childrens_book"):
     """Transform basic prompts into professional-grade prompts"""
     
@@ -309,6 +312,54 @@ def save_to_oci_via_api(image, prompt, story_title, page_number):
     except Exception as e:
         raise Exception(f"API upload failed: {str(e)}")
 
+def generate_test_image(prompt, model_choice, style_choice):
+    """Generate a single image for testing with better file management"""
+    try:
+        if not prompt.strip():
+            return None, "❌ Please enter a prompt"
+        
+        print(f"🎨 Generating test image with prompt: {prompt}")
+        
+        # Enhance the prompt
+        enhanced_prompt, negative_prompt = enhance_prompt(prompt, style_choice)
+        
+        # Generate the image
+        image = generate_high_quality_image(
+            enhanced_prompt, 
+            model_choice, 
+            style_choice,
+            negative_prompt
+        )
+        
+        # Save to accessible location
+        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"
+        filepath = f"generated_test_images/{filename}"
+        
+        try:
+            image.save(filepath)
+            save_info = f"💾 Saved to: {filepath}"
+            print(f"✅ Image saved: {filepath}")
+        except Exception as e:
+            save_info = f"⚠️ Could not save to disk: {e}"
+            print(f"⚠️ Save failed: {e}")
+        
+        status_msg = f"""✅ Success! Generated: {prompt}
+
+{save_info}
+
+🎨 Enhanced prompt: {enhanced_prompt}
+
+📥 Use the download button below the image to save it!"""
+        
+        return image, status_msg
+        
+    except Exception as e:
+        error_msg = f"❌ Generation failed: {str(e)}"
+        print(error_msg)
+        return None, error_msg
+
 # JOB MANAGEMENT FUNCTIONS
 def create_job(story_request: StorybookRequest) -> str:
     """Create a new job and return job ID"""
@@ -523,48 +574,10 @@ async def list_jobs():
         }
     }
 
-# SIMPLIFIED GRADIO INTERFACE FOR TESTING
+# GRADIO INTERFACE FOR TESTING
 def create_test_interface():
-    """Create a simple Gradio interface for testing image generation"""
+    """Create Gradio interface for testing image generation"""
     
-    def generate_test_image(prompt, model_choice, style_choice):
-        """Generate a single image for testing"""
-        try:
-            if not prompt.strip():
-                return None, "❌ Please enter a prompt"
-            
-            print(f"🎨 Generating test image with prompt: {prompt}")
-            
-            # Enhance the prompt
-            enhanced_prompt, negative_prompt = enhance_prompt(prompt, style_choice)
-            
-            # Generate the image
-            image = generate_high_quality_image(
-                enhanced_prompt, 
-                model_choice, 
-                style_choice,
-                negative_prompt
-            )
-            
-            # Create a simple save (optional)
-            try:
-                timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-                safe_prompt = "".join(c for c in prompt[:20] if c.isalnum() or c in (' ', '-', '_')).rstrip()
-                filename = f"test_{safe_prompt}_{timestamp}.png"
-                image.save(f"/tmp/{filename}")
-                print(f"💾 Test image saved as: {filename}")
-            except Exception as e:
-                print(f"⚠️ Could not save test image: {e}")
-            
-            status_msg = f"✅ Success! Generated: {prompt}\n\nEnhanced prompt: {enhanced_prompt}"
-            return image, status_msg
-            
-        except Exception as e:
-            error_msg = f"❌ Generation failed: {str(e)}"
-            print(error_msg)
-            return None, error_msg
-    
-    # Create the interface
     with gr.Blocks(title="High-Quality Image Generator", theme="soft") as demo:
         gr.Markdown("# 🎨 High-Quality Image Generator")
         gr.Markdown("Test your prompts with **studio-quality** image generation")
@@ -661,14 +674,17 @@ gradio_app = create_test_interface()
 # Mount Gradio interface to FastAPI
 @app.get("/")
 async def root():
-    return {"message": "Storybook Generator API is running! Visit /gradio for the testing interface."}
+    return {"message": "Storybook Generator API is running! Visit the main URL for the testing interface."}
 
-@app.get("/gradio")
-async def gradio_redirect():
-    """Redirect to Gradio interface"""
-    import urllib.parse
-    gradio_url = "http://localhost:7860" if __name__ == "__main__" else "/gradio"
-    return {"message": "Visit the Gradio interface at the main URL"}
+@app.get("/health")
+async def health_check():
+    """Health check endpoint"""
+    return {
+        "status": "healthy",
+        "models_loaded": list(model_cache.keys()),
+        "active_jobs": len(job_storage),
+        "memory_usage_mb": monitor_memory()
+    }
 
 # For Hugging Face Spaces deployment
 def get_app():
@@ -681,6 +697,7 @@ if __name__ == "__main__":
     print("🎨 Available models:")
     for model in MODEL_CHOICES:
         print(f"   - {model}")
+    print("📁 Images will be saved to: generated_test_images/")
     print("🌐 Web interface available at: http://localhost:7860")
     print("🔌 API endpoints available at: http://localhost:7860/api/")