Update app.py

app.py

@@ -24,7 +24,6 @@ import time
 from requests.adapters import HTTPAdapter
 from urllib3.util.retry import Retry
 from huggingface_hub import HfApi
-import accelerate
 import sys
 import traceback
 
@@ -38,7 +37,7 @@ 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
+# Simple diagnostic endpoint that works immediately
 app = FastAPI(title="Storybook Generator API")
 
 @app.get("/ping")
@@ -47,7 +46,7 @@ async def ping():
     return {
         "status": "alive",
         "timestamp": datetime.now().isoformat(),
-        "message": "Basic endpoint is working"
+        "message": "API is running"
     }
 
 @app.get("/debug")
@@ -73,22 +72,6 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# =============================================
-# MEMORY OPTIMIZATION SETTINGS
-# =============================================
-# Check if CUDA is available and set memory optimization
-if torch.cuda.is_available():
-    print("✅ CUDA available, enabling GPU optimizations")
-    torch.backends.cuda.enable_flash_sdp(True)
-    torch.backends.cuda.enable_mem_efficient_sdp(True)
-else:
-    print("⚠️ CUDA not available, running on CPU")
-
-# Set environment variables for memory optimization
-os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128"
-os.environ["OMP_NUM_THREADS"] = "1"
-os.environ["MKL_NUM_THREADS"] = "1"
-
 # =============================================
 # HUGGING FACE DATASET CONFIGURATION
 # =============================================
@@ -156,7 +139,7 @@ class MemoryStatusResponse(BaseModel):
     gpu_memory_cached_mb: Optional[float] = None
     status: str
 
-# HIGH-QUALITY MODEL SELECTION
+# HIGH-QUALITY MODEL SELECTION - SAME AS WORKING VERSION
 MODEL_CHOICES = {
     "dreamshaper-8": "lykon/dreamshaper-8",
     "realistic-vision": "SG161222/Realistic_Vision_V5.1", 
@@ -178,7 +161,7 @@ model_lock = threading.Lock()
 model_loading = False
 model_load_error = None
 
-# MEMORY MANAGEMENT FUNCTIONS
+# MEMORY MANAGEMENT FUNCTIONS - FROM WORKING VERSION
 def get_memory_usage():
     """Get current memory usage statistics"""
     process = psutil.Process()
@@ -267,10 +250,10 @@ def clear_memory(clear_models=True, clear_jobs=False, clear_local_images=False,
     }
 
 # =============================================
-# MODEL LOADING WITH PROPER ERROR HANDLING
+# SIMPLIFIED MODEL LOADING - EXACTLY LIKE WORKING VERSION
 # =============================================
 def load_model(model_name="dreamshaper-8"):
-    """Thread-safe model loading with error handling"""
+    """Thread-safe model loading - simplified like working version"""
     global model_cache, current_model_name, current_pipe, model_loading, model_load_error
     
     with model_lock:
@@ -279,10 +262,6 @@ 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
         
@@ -290,29 +269,23 @@ def load_model(model_name="dreamshaper-8"):
         try:
             model_id = MODEL_CHOICES.get(model_name, "lykon/dreamshaper-8")
             
-            # Load with memory optimizations
+            # Load model - exactly like your working version
             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,
-                use_safetensors=True
+                cache_dir="./model_cache"
             )
             
-            # Use memory efficient scheduler
+            # Use the same scheduler as working version
             pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
             
-            # Enable attention slicing
-            pipe.enable_attention_slicing()
-            
-            # Enable sequential CPU offload if needed
-            if not torch.cuda.is_available():
-                pipe.enable_sequential_cpu_offload()
-            
+            # Move to CPU - like working version
             pipe = pipe.to("cpu")
             
+            # NO additional optimizations - exactly like working version
+            
             model_cache[model_name] = pipe
             current_pipe = pipe
             current_model_name = model_name
@@ -332,12 +305,9 @@ def load_model(model_name="dreamshaper-8"):
                     "runwayml/stable-diffusion-v1-5", 
                     torch_dtype=torch.float32,
                     safety_checker=None,
-                    requires_safety_checker=False,
-                    low_cpu_mem_usage=True
+                    requires_safety_checker=False
                 ).to("cpu")
                 
-                pipe.enable_attention_slicing()
-                
                 model_cache[model_name] = pipe
                 current_pipe = pipe
                 current_model_name = "sd-1.5"
@@ -443,9 +413,9 @@ def upload_image_to_hf_dataset(image, project_id, page_number, prompt, style="")
         print(f"❌ Failed to upload image to HF Dataset: {e}")
         return None
 
-# PROMPT ENGINEERING
+# PROMPT ENGINEERING - FROM WORKING VERSION
 def enhance_prompt_simple(scene_visual, style="childrens_book"):
-    """Simple prompt enhancement"""
+    """Simple prompt enhancement - uses only the provided visual prompt with style"""
     
     style_templates = {
         "childrens_book": "children's book illustration, watercolor style, soft colors, whimsical, magical, storybook art, professional illustration",
@@ -466,13 +436,16 @@ def enhance_prompt_simple(scene_visual, style="childrens_book"):
     return enhanced_prompt, negative_prompt
 
 # =============================================
-# IMAGE GENERATION
+# IMAGE GENERATION - EXACTLY LIKE WORKING VERSION
 # =============================================
 def generate_image_simple(prompt, model_choice, style, scene_number, consistency_seed=None):
-    """Generate image with error handling"""
+    """Generate image - exactly like working version"""
     
     if current_pipe is None:
-        raise Exception("Model not loaded yet. Please wait a few seconds and try again.")
+        if model_loading:
+            raise Exception("Model is still loading. Please wait a few seconds and try again.")
+        else:
+            raise Exception(f"Model failed to load: {model_load_error}")
     
     enhanced_prompt, negative_prompt = enhance_prompt_simple(prompt, style)
     
@@ -484,19 +457,16 @@ def generate_image_simple(prompt, model_choice, style, scene_number, consistency
     try:
         pipe = current_pipe
         
-        with torch.inference_mode():
-            image = pipe(
-                prompt=enhanced_prompt,
-                negative_prompt=negative_prompt,
-                num_inference_steps=25,  # Reduced for speed
-                guidance_scale=7.5,
-                width=512,  # Reduced for memory
-                height=512,  # Reduced for memory
-                generator=torch.Generator(device="cpu").manual_seed(scene_seed)
-            ).images[0]
-        
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
+        # Use full quality settings like working version
+        image = pipe(
+            prompt=enhanced_prompt,
+            negative_prompt=negative_prompt,
+            num_inference_steps=35,
+            guidance_scale=7.5,
+            width=768,
+            height=1024,
+            generator=torch.Generator(device="cpu").manual_seed(scene_seed)
+        ).images[0]
         
         print(f"✅ Generated image for scene {scene_number}")
         return image
@@ -505,7 +475,7 @@ def generate_image_simple(prompt, model_choice, style, scene_number, consistency
         print(f"❌ Generation failed: {str(e)}")
         raise
 
-# LOCAL FILE MANAGEMENT FUNCTIONS
+# LOCAL FILE MANAGEMENT FUNCTIONS - FROM WORKING VERSION
 def save_image_to_local(image, prompt, style="test"):
     """Save image to local persistent storage"""
     try:
@@ -661,7 +631,7 @@ def calculate_remaining_time(job_id, progress):
 
 # BACKGROUND TASK
 def generate_storybook_background(job_id: str):
-    """Background task with memory optimization"""
+    """Background task to generate storybook"""
     try:
         if HF_TOKEN:
             ensure_dataset_exists()
@@ -790,12 +760,7 @@ async def root():
             "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\"}]}'"
-        }
+        "ui": "/ui"
     }
 
 @app.get("/api/health")
@@ -974,10 +939,11 @@ if __name__ == "__main__":
         print("📡 API endpoints:")
         print("   - GET  /ping")
         print("   - GET  /debug")
+        print("   - GET  /")
         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/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")