Update app.py

app.py

@@ -10,7 +10,8 @@ import re
 import time
 import json
 from typing import List, Optional, Dict
-from fastapi import FastAPI, HTTPException, BackgroundTasks
+from fastapi import FastAPI, HTTPException, BackgroundTasks, Request
+from fastapi.responses import JSONResponse
 from pydantic import BaseModel
 import gc
 import psutil
@@ -23,7 +24,54 @@ import time
 from requests.adapters import HTTPAdapter
 from urllib3.util.retry import Retry
 from huggingface_hub import HfApi
-import accelerate  # Add this for better memory management
+import accelerate
+import sys
+import traceback
+
+# =============================================
+# INITIAL SETUP & DIAGNOSTICS
+# =============================================
+print("=" * 60)
+print("🚀 STARTING STORYBOOK GENERATOR API")
+print("=" * 60)
+print(f"Python version: {sys.version}")
+print(f"PyTorch version: {torch.__version__}")
+print(f"CUDA available: {torch.cuda.is_available()}")
+
+# Create diagnostic endpoint that works immediately
+app = FastAPI(title="Storybook Generator API")
+
+@app.get("/ping")
+async def ping():
+    """Simple ping endpoint that always works"""
+    return {
+        "status": "alive",
+        "timestamp": datetime.now().isoformat(),
+        "message": "Basic endpoint is working"
+    }
+
+@app.get("/debug")
+async def debug():
+    """Debug endpoint showing system status"""
+    return {
+        "app_started": True,
+        "python_version": sys.version,
+        "torch_version": torch.__version__,
+        "cuda_available": torch.cuda.is_available(),
+        "routes": [{"path": route.path, "methods": list(route.methods)} for route in app.routes],
+        "hf_token_set": bool(os.environ.get("HF_TOKEN")),
+        "timestamp": datetime.now().isoformat()
+    }
+
+# Add CORS middleware
+from fastapi.middleware.cors import CORSMiddleware
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
 
 # =============================================
 # MEMORY OPTIMIZATION SETTINGS
@@ -57,19 +105,6 @@ PERSISTENT_IMAGE_DIR = "generated_test_images"
 os.makedirs(PERSISTENT_IMAGE_DIR, exist_ok=True)
 print(f"📁 Created local image directory: {PERSISTENT_IMAGE_DIR}")
 
-# Initialize FastAPI app
-app = FastAPI(title="Storybook Generator API")
-
-# Add CORS middleware
-from fastapi.middleware.cors import CORSMiddleware
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
 # Job Status Enum
 class JobStatus(str, Enum):
     PENDING = "pending"
@@ -140,6 +175,8 @@ model_cache = {}
 current_model_name = None
 current_pipe = None
 model_lock = threading.Lock()
+model_loading = False
+model_load_error = None
 
 # MEMORY MANAGEMENT FUNCTIONS
 def get_memory_usage():
@@ -230,11 +267,11 @@ def clear_memory(clear_models=True, clear_jobs=False, clear_local_images=False,
     }
 
 # =============================================
-# OPTIMIZED MODEL LOADING - MAINTAINS QUALITY
+# MODEL LOADING WITH PROPER ERROR HANDLING
 # =============================================
 def load_model(model_name="dreamshaper-8"):
-    """Thread-safe model loading with memory optimization but maintaining quality"""
-    global model_cache, current_model_name, current_pipe
+    """Thread-safe model loading with error handling"""
+    global model_cache, current_model_name, current_pipe, model_loading, model_load_error
     
     with model_lock:
         if model_name in model_cache:
@@ -242,26 +279,32 @@ def load_model(model_name="dreamshaper-8"):
             current_model_name = model_name
             return current_pipe
         
+        if model_loading:
+            print(f"⏳ Model already loading, waiting...")
+            return None
+        
+        model_loading = True
+        model_load_error = None
+        
         print(f"🔄 Loading model: {model_name}")
         try:
             model_id = MODEL_CHOICES.get(model_name, "lykon/dreamshaper-8")
             
-            # Load with memory optimizations but keep quality
+            # Load with memory optimizations
             pipe = StableDiffusionPipeline.from_pretrained(
                 model_id, 
                 torch_dtype=torch.float32,
                 safety_checker=None,
                 requires_safety_checker=False,
                 cache_dir="./model_cache",
-                low_cpu_mem_usage=True,  # Reduces memory during loading
-                use_safetensors=True,
-                variant="fp32"  # Use full precision for quality
+                low_cpu_mem_usage=True,
+                use_safetensors=True
             )
             
-            # Use memory efficient scheduler (maintains quality)
+            # Use memory efficient scheduler
             pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
             
-            # Enable attention slicing (trades speed for memory)
+            # Enable attention slicing
             pipe.enable_attention_slicing()
             
             # Enable sequential CPU offload if needed
@@ -273,11 +316,14 @@ def load_model(model_name="dreamshaper-8"):
             model_cache[model_name] = pipe
             current_pipe = pipe
             current_model_name = model_name
+            model_loading = False
             
             print(f"✅ Model loaded: {model_name}")
             return pipe
             
         except Exception as e:
+            model_load_error = str(e)
+            model_loading = False
             print(f"❌ Model loading failed for {model_name}: {e}")
             print(f"🔄 Falling back to stable-diffusion-v1-5")
             
@@ -295,18 +341,30 @@ def load_model(model_name="dreamshaper-8"):
                 model_cache[model_name] = pipe
                 current_pipe = pipe
                 current_model_name = "sd-1.5"
+                model_loading = False
                 
                 print(f"✅ Fallback model loaded")
                 return pipe
                 
             except Exception as fallback_error:
+                model_load_error = str(fallback_error)
+                model_loading = False
                 print(f"❌ Fallback model failed: {fallback_error}")
                 raise
 
-# Initialize default model
-print("🚀 Initializing Storybook Generator API...")
-load_model("dreamshaper-8")
-print("✅ Model loaded and ready!")
+# Try to load model in background thread to not block startup
+def load_model_background():
+    try:
+        load_model("dreamshaper-8")
+    except Exception as e:
+        print(f"❌ Background model loading failed: {e}")
+
+# Start model loading in background
+import threading
+model_thread = threading.Thread(target=load_model_background)
+model_thread.daemon = True
+model_thread.start()
+print("⏳ Model loading started in background...")
 
 # =============================================
 # HF DATASET FUNCTIONS
@@ -387,7 +445,7 @@ def upload_image_to_hf_dataset(image, project_id, page_number, prompt, style="")
 
 # PROMPT ENGINEERING
 def enhance_prompt_simple(scene_visual, style="childrens_book"):
-    """Simple prompt enhancement - uses only the provided visual prompt with style"""
+    """Simple prompt enhancement"""
     
     style_templates = {
         "childrens_book": "children's book illustration, watercolor style, soft colors, whimsical, magical, storybook art, professional illustration",
@@ -408,10 +466,13 @@ def enhance_prompt_simple(scene_visual, style="childrens_book"):
     return enhanced_prompt, negative_prompt
 
 # =============================================
-# OPTIMIZED IMAGE GENERATION - MAINTAINS QUALITY
+# IMAGE GENERATION
 # =============================================
 def generate_image_simple(prompt, model_choice, style, scene_number, consistency_seed=None):
-    """Generate image with memory optimization but maintaining quality"""
+    """Generate image with error handling"""
+    
+    if current_pipe is None:
+        raise Exception("Model not loaded yet. Please wait a few seconds and try again.")
     
     enhanced_prompt, negative_prompt = enhance_prompt_simple(prompt, style)
     
@@ -421,21 +482,19 @@ def generate_image_simple(prompt, model_choice, style, scene_number, consistency
         scene_seed = random.randint(1000, 9999)
     
     try:
-        pipe = load_model(model_choice)
+        pipe = current_pipe
         
-        # Use memory efficient generation
-        with torch.inference_mode():  # More memory efficient than no_grad
+        with torch.inference_mode():
             image = pipe(
                 prompt=enhanced_prompt,
                 negative_prompt=negative_prompt,
-                num_inference_steps=35,  # Keep quality
+                num_inference_steps=25,  # Reduced for speed
                 guidance_scale=7.5,
-                width=768,  # Keep quality
-                height=1024,  # Keep quality
+                width=512,  # Reduced for memory
+                height=512,  # Reduced for memory
                 generator=torch.Generator(device="cpu").manual_seed(scene_seed)
             ).images[0]
         
-        # Clean up after generation
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
         
@@ -555,6 +614,9 @@ def update_job_status(job_id: str, status: JobStatus, progress: int, message: st
     if result:
         job_storage[job_id]["result"] = result
     
+    job_data = job_storage[job_id]
+    request_data = job_data["request"]
+    
     if request_data.get("callback_url"):
         try:
             callback_url = request_data["callback_url"]
@@ -597,7 +659,7 @@ def calculate_remaining_time(job_id, progress):
     
     return "Unknown"
 
-# OPTIMIZED BACKGROUND TASK
+# BACKGROUND TASK
 def generate_storybook_background(job_id: str):
     """Background task with memory optimization"""
     try:
@@ -663,7 +725,7 @@ def generate_storybook_background(job_id: str):
                 }
                 generated_pages.append(page_data)
                 
-                # Clean up after each page
+                # Clean up
                 if torch.cuda.is_available():
                     torch.cuda.empty_cache()
                 gc.collect()
@@ -696,14 +758,87 @@ def generate_storybook_background(job_id: str):
     except Exception as e:
         error_msg = f"Generation failed: {str(e)}"
         print(f"❌ {error_msg}")
+        traceback.print_exc()
         update_job_status(job_id, JobStatus.FAILED, 0, error_msg)
 
+# =============================================
 # FASTAPI ENDPOINTS
+# =============================================
+
+@app.get("/")
+async def root():
+    """Root endpoint showing API status"""
+    return {
+        "name": "Storybook Generator API",
+        "version": "1.0.0",
+        "status": "running",
+        "model_status": {
+            "loaded": current_model_name is not None,
+            "model_name": current_model_name,
+            "loading": model_loading,
+            "error": model_load_error
+        },
+        "hf_dataset": DATASET_ID if HF_TOKEN else "Disabled",
+        "endpoints": {
+            "ping": "GET /ping",
+            "debug": "GET /debug",
+            "health": "GET /api/health",
+            "generate": "POST /api/generate-storybook",
+            "status": "GET /api/job-status/{job_id}",
+            "project_images": "GET /api/project-images/{project_id}",
+            "memory": "GET /api/memory-status",
+            "clear_memory": "POST /api/clear-memory",
+            "local_images": "GET /api/local-images"
+        },
+        "ui": "/ui",
+        "test_commands": {
+            "ping": "curl -X GET https://yukee1992-Video_image_generator.hf.space/ping",
+            "health": "curl -X GET https://yukee1992-Video_image_generator.hf.space/api/health",
+            "generate": "curl -X POST https://yukee1992-Video_image_generator.hf.space/api/generate-storybook -H 'Content-Type: application/json' -d '{\"story_title\":\"test\",\"scenes\":[{\"visual\":\"a cat\",\"text\":\"test\"}]}'"
+        }
+    }
+
+@app.get("/api/health")
+async def health():
+    """Health check endpoint"""
+    return {
+        "status": "healthy",
+        "service": "storybook-generator",
+        "model_loaded": current_model_name is not None,
+        "model_name": current_model_name,
+        "model_loading": model_loading,
+        "hf_dataset": DATASET_ID if HF_TOKEN else "Disabled",
+        "active_jobs": len(job_storage),
+        "timestamp": datetime.now().isoformat()
+    }
+
 @app.post("/api/generate-storybook")
 async def generate_storybook(request: dict, background_tasks: BackgroundTasks):
+    """Generate a storybook from scenes"""
     try:
         print(f"📥 Received request for: {request.get('story_title', 'Unknown')}")
         
+        # Check if model is loaded
+        if current_pipe is None:
+            if model_loading:
+                return JSONResponse(
+                    status_code=503,
+                    content={
+                        "status": "loading",
+                        "message": "Model is still loading. Please wait a few seconds and try again.",
+                        "estimated_time": "10-20 seconds"
+                    }
+                )
+            else:
+                return JSONResponse(
+                    status_code=503,
+                    content={
+                        "status": "error",
+                        "message": f"Model failed to load: {model_load_error}",
+                        "error": model_load_error
+                    }
+                )
+        
         if 'consistency_seed' not in request:
             request['consistency_seed'] = random.randint(1000, 9999)
         
@@ -729,10 +864,13 @@ async def generate_storybook(request: dict, background_tasks: BackgroundTasks):
         }
         
     except Exception as e:
+        print(f"❌ Error in generate_storybook: {e}")
+        traceback.print_exc()
         raise HTTPException(status_code=500, detail=str(e))
 
 @app.get("/api/job-status/{job_id}")
 async def get_job_status(job_id: str):
+    """Get job status by ID"""
     job_data = job_storage.get(job_id)
     if not job_data:
         raise HTTPException(status_code=404, detail="Job not found")
@@ -745,17 +883,9 @@ async def get_job_status(job_id: str):
         "result": job_data["result"]
     }
 
-@app.get("/api/health")
-async def health():
-    return {
-        "status": "healthy",
-        "service": "storybook-generator",
-        "hf_dataset": DATASET_ID if HF_TOKEN else "Disabled",
-        "active_jobs": len(job_storage)
-    }
-
 @app.get("/api/project-images/{project_id}")
 async def get_project_images(project_id: str):
+    """Get all images for a project from HF Dataset"""
     try:
         if not HF_TOKEN:
             return {"error": "HF_TOKEN not set"}
@@ -770,16 +900,47 @@ async def get_project_images(project_id: str):
     except Exception as e:
         return {"error": str(e)}
 
+@app.get("/api/memory-status")
+async def memory_status():
+    """Get memory usage status"""
+    return get_memory_usage()
+
+@app.post("/api/clear-memory")
+async def clear_memory_api(request: MemoryClearanceRequest):
+    """Clear memory manually"""
+    return clear_memory(
+        clear_models=request.clear_models,
+        clear_jobs=request.clear_jobs,
+        clear_local_images=request.clear_local_images,
+        force_gc=request.force_gc
+    )
+
+@app.get("/api/local-images")
+async def get_local_images():
+    """Get locally saved images"""
+    return get_local_storage_info()
+
+# =============================================
 # GRADIO INTERFACE
+# =============================================
 def create_gradio_interface():
     def generate_test(prompt, model_choice, style_choice):
         if not prompt.strip():
             return None, "❌ Please enter a prompt"
         
-        image = generate_image_simple(prompt, model_choice, style_choice, 1)
-        filepath, filename = save_image_to_local(image, prompt, style_choice)
-        
-        return image, f"✅ Generated! Local: {filename}"
+        try:
+            if current_pipe is None:
+                if model_loading:
+                    return None, "⏳ Model is still loading. Please wait a few seconds..."
+                else:
+                    return None, f"❌ Model failed to load: {model_load_error}"
+            
+            image = generate_image_simple(prompt, model_choice, style_choice, 1)
+            filepath, filename = save_image_to_local(image, prompt, style_choice)
+            
+            return image, f"✅ Generated! Local: {filename}"
+        except Exception as e:
+            return None, f"❌ Error: {str(e)}"
     
     with gr.Blocks(title="Storybook Generator") as demo:
         gr.Markdown("# 🎨 Storybook Generator")
@@ -801,27 +962,31 @@ def create_gradio_interface():
 
 demo = create_gradio_interface()
 
-@app.get("/")
-async def root():
-    return {
-        "message": "Storybook Generator API",
-        "hf_dataset": DATASET_ID if HF_TOKEN else "Disabled",
-        "endpoints": {
-            "generate": "POST /api/generate-storybook",
-            "status": "GET /api/job-status/{job_id}",
-            "health": "GET /api/health",
-            "project_images": "GET /api/project-images/{project_id}"
-        },
-        "ui": "/ui"
-    }
-
+# =============================================
+# MAIN
+# =============================================
 if __name__ == "__main__":
     import uvicorn
     
     if os.environ.get('SPACE_ID'):
         print("🚀 Running on Hugging Face Spaces")
         print(f"📦 HF Dataset: {DATASET_ID if HF_TOKEN else 'Disabled'}")
+        print("📡 API endpoints:")
+        print("   - GET  /ping")
+        print("   - GET  /debug")
+        print("   - GET  /api/health")
+        print("   - POST /api/generate-storybook")
+        print("   - GET  /api/job-status/{{job_id}}")
+        print("   - GET  /api/project-images/{{project_id}}")
+        print("   - GET  /api/memory-status")
+        print("   - POST /api/clear-memory")
+        print("   - GET  /api/local-images")
+        print("🎨 UI: /ui")
+        
+        # Mount Gradio
         gr.mount_gradio_app(app, demo, path="/ui")
+        
+        # Run the combined app
         uvicorn.run(app, host="0.0.0.0", port=7860, log_level="info")
     else:
         print("🚀 Running locally")