Update app.py

app.py

@@ -22,10 +22,24 @@ import random
 import time
 from requests.adapters import HTTPAdapter
 from urllib3.util.retry import Retry
-from huggingface_hub import HfApi  # NEW: Add this import
+from huggingface_hub import HfApi
+import accelerate  # Add this for better memory management
 
 # =============================================
-# HUGGING FACE DATASET CONFIGURATION (NEW)
+# MEMORY OPTIMIZATION SETTINGS
+# =============================================
+# Enable memory efficient attention if available
+if hasattr(torch, 'backends') and hasattr(torch.backends, 'cuda') and torch.backends.cuda.is_enabled():
+    torch.backends.cuda.enable_flash_sdp(True)
+    torch.backends.cuda.enable_mem_efficient_sdp(True)
+
+# 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
 # =============================================
 HF_TOKEN = os.environ.get("HF_TOKEN")
 HF_USERNAME = "yukee1992"
@@ -77,7 +91,7 @@ class StorybookRequest(BaseModel):
     style: str = "childrens_book"
     callback_url: Optional[str] = None
     consistency_seed: Optional[int] = None
-    project_id: Optional[str] = None  # ADDED for HF Dataset organization
+    project_id: Optional[str] = None
 
 class JobStatusResponse(BaseModel):
     job_id: str
@@ -104,7 +118,7 @@ class MemoryStatusResponse(BaseModel):
     gpu_memory_cached_mb: Optional[float] = None
     status: str
 
-# HIGH-QUALITY MODEL SELECTION - ANIME FOCUSED & WORKING
+# HIGH-QUALITY MODEL SELECTION
 MODEL_CHOICES = {
     "dreamshaper-8": "lykon/dreamshaper-8",
     "realistic-vision": "SG161222/Realistic_Vision_V5.1", 
@@ -132,7 +146,6 @@ def get_memory_usage():
     memory_used_mb = memory_info.rss / (1024 * 1024)
     memory_percent = process.memory_percent()
     
-    # GPU memory if available
     gpu_memory_allocated_mb = None
     gpu_memory_cached_mb = None
     
@@ -154,17 +167,13 @@ def clear_memory(clear_models=True, clear_jobs=False, clear_local_images=False,
     """Clear memory by unloading models and cleaning up resources"""
     results = []
     
-    # Clear model cache
     if clear_models:
         with model_lock:
             models_cleared = len(model_cache)
             for model_name, pipe in model_cache.items():
                 try:
-                    # Move to CPU first if it's on GPU
                     if hasattr(pipe, 'to'):
                         pipe.to('cpu')
-                    
-                    # Delete the pipeline
                     del pipe
                     results.append(f"Unloaded model: {model_name}")
                 except Exception as e:
@@ -176,7 +185,6 @@ def clear_memory(clear_models=True, clear_jobs=False, clear_local_images=False,
             current_model_name = None
             results.append(f"Cleared {models_cleared} models from cache")
     
-    # Clear completed jobs
     if clear_jobs:
         jobs_to_clear = []
         for job_id, job_data in job_storage.items():
@@ -189,7 +197,6 @@ def clear_memory(clear_models=True, clear_jobs=False, clear_local_images=False,
         
         results.append(f"Cleared {len(jobs_to_clear)} completed/failed jobs")
     
-    # Clear local images
     if clear_local_images:
         try:
             storage_info = get_local_storage_info()
@@ -203,7 +210,6 @@ def clear_memory(clear_models=True, clear_jobs=False, clear_local_images=False,
         except Exception as e:
             results.append(f"Error clearing local images: {str(e)}")
     
-    # Force garbage collection
     if force_gc:
         gc.collect()
         if torch.cuda.is_available():
@@ -212,7 +218,6 @@ def clear_memory(clear_models=True, clear_jobs=False, clear_local_images=False,
             results.append("GPU cache cleared")
         results.append("Garbage collection forced")
     
-    # Get memory status after cleanup
     memory_status = get_memory_usage()
     
     return {
@@ -221,8 +226,11 @@ def clear_memory(clear_models=True, clear_jobs=False, clear_local_images=False,
         "memory_after_cleanup": memory_status
     }
 
+# =============================================
+# OPTIMIZED MODEL LOADING - MAINTAINS QUALITY
+# =============================================
 def load_model(model_name="dreamshaper-8"):
-    """Thread-safe model loading with HIGH-QUALITY settings and better error handling"""
+    """Thread-safe model loading with memory optimization but maintaining quality"""
     global model_cache, current_model_name, current_pipe
     
     with model_lock:
@@ -231,53 +239,65 @@ def load_model(model_name="dreamshaper-8"):
             current_model_name = model_name
             return current_pipe
         
-        print(f"🔄 Loading HIGH-QUALITY model: {model_name}")
+        print(f"🔄 Loading model: {model_name}")
         try:
             model_id = MODEL_CHOICES.get(model_name, "lykon/dreamshaper-8")
             
-            print(f"🔧 Attempting to load: {model_id}")
-            
+            # Load with memory optimizations but keep quality
             pipe = StableDiffusionPipeline.from_pretrained(
                 model_id, 
                 torch_dtype=torch.float32,
                 safety_checker=None,
                 requires_safety_checker=False,
-                local_files_only=False,  # Allow downloading if not cached
-                cache_dir="./model_cache"  # Specific cache directory
+                cache_dir="./model_cache",
+                low_cpu_mem_usage=True,  # Reduces memory during loading
+                use_safetensors=True,
+                variant="fp32"  # Use full precision for quality
             )
             
+            # Use memory efficient scheduler (maintains quality)
             pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
+            
+            # Enable attention slicing (trades speed for memory)
+            pipe.enable_attention_slicing()
+            
+            # Enable sequential CPU offload if needed
+            if not torch.cuda.is_available():
+                pipe.enable_sequential_cpu_offload()
+            
             pipe = pipe.to("cpu")
             
             model_cache[model_name] = pipe
             current_pipe = pipe
             current_model_name = model_name
             
-            print(f"✅ HIGH-QUALITY Model loaded: {model_name}")
+            print(f"✅ Model loaded: {model_name}")
             return pipe
             
         except Exception as e:
             print(f"❌ Model loading failed for {model_name}: {e}")
             print(f"🔄 Falling back to stable-diffusion-v1-5")
             
-            # Fallback to base model
             try:
                 pipe = StableDiffusionPipeline.from_pretrained(
                     "runwayml/stable-diffusion-v1-5", 
                     torch_dtype=torch.float32,
                     safety_checker=None,
-                    requires_safety_checker=False
+                    requires_safety_checker=False,
+                    low_cpu_mem_usage=True
                 ).to("cpu")
                 
+                pipe.enable_attention_slicing()
+                
                 model_cache[model_name] = pipe
                 current_pipe = pipe
                 current_model_name = "sd-1.5"
                 
-                print(f"✅ Fallback model loaded: stable-diffusion-v1-5")
+                print(f"✅ Fallback model loaded")
                 return pipe
                 
             except Exception as fallback_error:
-                print(f"❌ Critical: Fallback model also failed: {fallback_error}")
+                print(f"❌ Fallback model failed: {fallback_error}")
                 raise
 
 # Initialize default model
@@ -285,11 +305,87 @@ print("🚀 Initializing Storybook Generator API...")
 load_model("dreamshaper-8")
 print("✅ Model loaded and ready!")
 
-# SIMPLE PROMPT ENGINEERING - USE PURE PROMPTS ONLY
+# =============================================
+# HF DATASET FUNCTIONS
+# =============================================
+def ensure_dataset_exists():
+    """Create dataset if it doesn't exist"""
+    if not HF_TOKEN:
+        print("⚠️ HF_TOKEN not set, cannot create/verify dataset")
+        return False
+    
+    try:
+        api = HfApi(token=HF_TOKEN)
+        try:
+            api.dataset_info(DATASET_ID)
+            print(f"✅ Dataset {DATASET_ID} exists")
+        except Exception:
+            print(f"📦 Creating dataset: {DATASET_ID}")
+            api.create_repo(
+                repo_id=DATASET_ID,
+                repo_type="dataset",
+                private=False,
+                exist_ok=True
+            )
+            print(f"✅ Created dataset: {DATASET_ID}")
+        return True
+    except Exception as e:
+        print(f"❌ Failed to ensure dataset: {e}")
+        return False
+
+def upload_to_hf_dataset(file_content, filename, subfolder=""):
+    """Upload a file to Hugging Face Dataset"""
+    if not HF_TOKEN:
+        print("⚠️ HF_TOKEN not set, skipping upload")
+        return None
+    
+    try:
+        if subfolder:
+            path_in_repo = f"data/{subfolder}/{filename}"
+        else:
+            path_in_repo = f"data/{filename}"
+        
+        api = HfApi(token=HF_TOKEN)
+        api.upload_file(
+            path_or_fileobj=file_content,
+            path_in_repo=path_in_repo,
+            repo_id=DATASET_ID,
+            repo_type="dataset"
+        )
+        
+        url = f"https://huggingface.co/datasets/{DATASET_ID}/resolve/main/{path_in_repo}"
+        print(f"✅ Uploaded to HF Dataset: {url}")
+        return url
+        
+    except Exception as e:
+        print(f"❌ Failed to upload to HF Dataset: {e}")
+        return None
+
+def upload_image_to_hf_dataset(image, project_id, page_number, prompt, style=""):
+    """Upload generated image to HF Dataset"""
+    try:
+        img_bytes = io.BytesIO()
+        image.save(img_bytes, format='PNG')
+        img_data = img_bytes.getvalue()
+        
+        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()
+        safe_prompt = safe_prompt.replace(' ', '_')
+        filename = f"page_{page_number:03d}_{safe_prompt}_{timestamp}.png"
+        
+        subfolder = f"projects/{project_id}"
+        url = upload_to_hf_dataset(img_data, filename, subfolder)
+        
+        return url
+        
+    except Exception as e:
+        print(f"❌ Failed to upload image to HF Dataset: {e}")
+        return None
+
+# PROMPT ENGINEERING
 def enhance_prompt_simple(scene_visual, style="childrens_book"):
     """Simple prompt enhancement - uses only the provided visual prompt with style"""
     
-    # Style templates
     style_templates = {
         "childrens_book": "children's book illustration, watercolor style, soft colors, whimsical, magical, storybook art, professional illustration",
         "realistic": "photorealistic, detailed, natural lighting, professional photography",
@@ -299,10 +395,8 @@ def enhance_prompt_simple(scene_visual, style="childrens_book"):
     
     style_prompt = style_templates.get(style, style_templates["childrens_book"])
     
-    # Use only the provided visual prompt with style
     enhanced_prompt = f"{style_prompt}, {scene_visual}"
     
-    # Basic negative prompt for quality
     negative_prompt = (
         "blurry, low quality, bad anatomy, deformed characters, "
         "wrong proportions, mismatched features"
@@ -310,13 +404,14 @@ def enhance_prompt_simple(scene_visual, style="childrens_book"):
     
     return enhanced_prompt, negative_prompt
 
+# =============================================
+# OPTIMIZED IMAGE GENERATION - MAINTAINS QUALITY
+# =============================================
 def generate_image_simple(prompt, model_choice, style, scene_number, consistency_seed=None):
-    """Generate image using pure prompts only"""
+    """Generate image with memory optimization but maintaining quality"""
     
-    # Enhance prompt with simple style addition
     enhanced_prompt, negative_prompt = enhance_prompt_simple(prompt, style)
     
-    # Use seed if provided
     if consistency_seed:
         scene_seed = consistency_seed + scene_number
     else:
@@ -325,20 +420,23 @@ def generate_image_simple(prompt, model_choice, style, scene_number, consistency
     try:
         pipe = load_model(model_choice)
         
-        image = pipe(
-            prompt=enhanced_prompt,
-            negative_prompt=negative_prompt,
-            num_inference_steps=35,
-            guidance_scale=7.5,
-            width=768,
-            height=1024,  # Portrait for better full-body
-            generator=torch.Generator(device="cpu").manual_seed(scene_seed)
-        ).images[0]
+        # Use memory efficient generation
+        with torch.inference_mode():  # More memory efficient than no_grad
+            image = pipe(
+                prompt=enhanced_prompt,
+                negative_prompt=negative_prompt,
+                num_inference_steps=35,  # Keep quality
+                guidance_scale=7.5,
+                width=768,  # Keep quality
+                height=1024,  # Keep quality
+                generator=torch.Generator(device="cpu").manual_seed(scene_seed)
+            ).images[0]
+        
+        # Clean up after generation
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
         
         print(f"✅ Generated image for scene {scene_number}")
-        print(f"🌱 Seed used: {scene_seed}")
-        print(f"📝 Pure prompt used: {prompt}")
-        
         return image
         
     except Exception as e:
@@ -353,12 +451,10 @@ def save_image_to_local(image, prompt, style="test"):
         safe_prompt = "".join(c for c in prompt[:50] if c.isalnum() or c in (' ', '-', '_')).rstrip()
         filename = f"image_{safe_prompt}_{timestamp}.png"
         
-        # Create style subfolder
         style_dir = os.path.join(PERSISTENT_IMAGE_DIR, style)
         os.makedirs(style_dir, exist_ok=True)
         filepath = os.path.join(style_dir, filename)
         
-        # Save the image
         image.save(filepath)
         print(f"💾 Image saved locally: {filepath}")
         
@@ -373,7 +469,6 @@ def delete_local_image(filepath):
     try:
         if os.path.exists(filepath):
             os.remove(filepath)
-            print(f"🗑️ Deleted local image: {filepath}")
             return True, f"✅ Deleted: {os.path.basename(filepath)}"
         else:
             return False, f"❌ File not found: {filepath}"
@@ -425,89 +520,6 @@ def refresh_local_images():
         print(f"Error refreshing local images: {e}")
         return []
 
-# =============================================
-# NEW: HUGGING FACE DATASET FUNCTIONS
-# =============================================
-
-def ensure_dataset_exists():
-    """Create dataset if it doesn't exist"""
-    if not HF_TOKEN:
-        print("⚠️ HF_TOKEN not set, cannot create/verify dataset")
-        return False
-    
-    try:
-        api = HfApi(token=HF_TOKEN)
-        try:
-            api.dataset_info(DATASET_ID)
-            print(f"✅ Dataset {DATASET_ID} exists")
-        except Exception:
-            print(f"📦 Creating dataset: {DATASET_ID}")
-            api.create_repo(
-                repo_id=DATASET_ID,
-                repo_type="dataset",
-                private=False,
-                exist_ok=True
-            )
-            print(f"✅ Created dataset: {DATASET_ID}")
-        return True
-    except Exception as e:
-        print(f"❌ Failed to ensure dataset: {e}")
-        return False
-
-def upload_to_hf_dataset(file_content, filename, subfolder=""):
-    """Upload a file to Hugging Face Dataset"""
-    if not HF_TOKEN:
-        print("⚠️ HF_TOKEN not set, skipping upload")
-        return None
-    
-    try:
-        if subfolder:
-            path_in_repo = f"data/{subfolder}/{filename}"
-        else:
-            path_in_repo = f"data/{filename}"
-        
-        api = HfApi(token=HF_TOKEN)
-        api.upload_file(
-            path_or_fileobj=file_content,
-            path_in_repo=path_in_repo,
-            repo_id=DATASET_ID,
-            repo_type="dataset"
-        )
-        
-        url = f"https://huggingface.co/datasets/{DATASET_ID}/resolve/main/{path_in_repo}"
-        print(f"✅ Uploaded to HF Dataset: {url}")
-        return url
-        
-    except Exception as e:
-        print(f"❌ Failed to upload to HF Dataset: {e}")
-        return None
-
-def upload_image_to_hf_dataset(image, project_id, page_number, prompt, style=""):
-    """Upload generated image to HF Dataset"""
-    try:
-        img_bytes = io.BytesIO()
-        image.save(img_bytes, format='PNG')
-        img_data = img_bytes.getvalue()
-        
-        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()
-        safe_prompt = safe_prompt.replace(' ', '_')
-        filename = f"page_{page_number:03d}_{safe_prompt}_{timestamp}.png"
-        
-        subfolder = f"projects/{project_id}"
-        url = upload_to_hf_dataset(img_data, filename, subfolder)
-        
-        return url
-        
-    except Exception as e:
-        print(f"❌ Failed to upload image to HF Dataset: {e}")
-        return None
-
-# =============================================
-# REMOVED: OCI BUCKET FUNCTIONS
-# (save_to_oci_bucket and test_oci_connection are removed)
-# =============================================
-
 # JOB MANAGEMENT FUNCTIONS
 def create_job(story_request: StorybookRequest) -> str:
     job_id = str(uuid.uuid4())
@@ -524,8 +536,6 @@ def create_job(story_request: StorybookRequest) -> str:
     }
     
     print(f"📝 Created job {job_id} for story: {story_request.story_title}")
-    print(f"📄 Scenes to generate: {len(story_request.scenes)}")
-    
     return job_id
 
 def update_job_status(job_id: str, status: JobStatus, progress: int, message: str, result=None):
@@ -542,66 +552,28 @@ def update_job_status(job_id: str, status: JobStatus, progress: int, message: st
     if result:
         job_storage[job_id]["result"] = result
     
-    # Send webhook notification if callback URL exists
-    job_data = job_storage[job_id]
-    request_data = job_data["request"]
-    
     if request_data.get("callback_url"):
         try:
             callback_url = request_data["callback_url"]
-            
             callback_data = {
                 "job_id": job_id,
                 "status": status.value,
                 "progress": progress,
                 "message": message,
                 "story_title": request_data["story_title"],
-                "total_scenes": len(request_data["scenes"]),
-                "timestamp": time.time(),
-                "source": "huggingface-image-generator",
-                "estimated_time_remaining": calculate_remaining_time(job_id, progress)
+                "timestamp": time.time()
             }
             
-            if status == JobStatus.PROCESSING:
-                total_scenes = len(request_data["scenes"])
-                if total_scenes > 0:
-                    current_scene = min((progress - 5) // (90 // total_scenes) + 1, total_scenes)
-                    callback_data["current_scene"] = current_scene
-                    callback_data["total_scenes"] = total_scenes
-                    
-                    if current_scene <= len(request_data["scenes"]):
-                        scene_data = request_data["scenes"][current_scene-1]
-                        callback_data["scene_description"] = scene_data.get("visual", "")[:100] + "..."
-                        callback_data["current_prompt"] = scene_data.get("visual", "")
-            
             if status == JobStatus.COMPLETED and result:
                 callback_data["result"] = {
-                    "total_pages": result.get("total_pages", 0),
-                    "generation_time": result.get("generation_time", 0),
-                    "hf_dataset_url": result.get("hf_dataset_url", ""),
-                    "pages_generated": result.get("generated_pages", 0),
-                    "consistency_seed": result.get("consistency_seed", None),
-                    "image_urls": result.get("image_urls", [])
+                    "image_urls": result.get("image_urls", []),
+                    "project_id": result.get("project_id", "")
                 }
             
-            headers = {
-                'Content-Type': 'application/json',
-                'User-Agent': 'Storybook-Generator/1.0'
-            }
-            
-            print(f"📢 Sending callback to: {callback_url}")
-            
-            response = requests.post(
-                callback_url,
-                json=callback_data,
-                headers=headers,
-                timeout=30
-            )
-            
-            print(f"📢 Callback sent: Status {response.status_code}")
-                
+            requests.post(callback_url, json=callback_data, timeout=5)
+            print(f"📢 Callback sent to {callback_url}")
         except Exception as e:
-            print(f"⚠️ Callback failed: {str(e)}")
+            print(f"⚠️ Callback failed: {e}")
     
     return True
 
@@ -622,28 +594,21 @@ def calculate_remaining_time(job_id, progress):
     
     return "Unknown"
 
-# UPDATED BACKGROUND TASK - Uses HF Dataset instead of OCI
+# OPTIMIZED BACKGROUND TASK
 def generate_storybook_background(job_id: str):
-    """Background task to generate complete storybook and upload to HF Dataset"""
+    """Background task with memory optimization"""
     try:
-        # Ensure HF Dataset exists
         if HF_TOKEN:
             ensure_dataset_exists()
         
         job_data = job_storage[job_id]
-        story_request_data = job_data["request"]
-        story_request = StorybookRequest(**story_request_data)
+        story_request = StorybookRequest(**job_data["request"])
         
-        # Use project_id from request or generate from story title
         project_id = story_request.project_id or story_request.story_title.replace(' ', '_').lower()
         
         print(f"🎬 Starting storybook generation for job {job_id}")
-        print(f"📖 Story: {story_request.story_title}")
-        print(f"📄 Scenes: {len(story_request.scenes)}")
-        print(f"🎨 Style: {story_request.style}")
-        print(f"📦 Project ID: {project_id}")
         
-        update_job_status(job_id, JobStatus.PROCESSING, 5, "Starting storybook generation with pure prompts...")
+        update_job_status(job_id, JobStatus.PROCESSING, 5, "Starting generation...")
         
         total_scenes = len(story_request.scenes)
         generated_pages = []
@@ -657,14 +622,11 @@ def generate_storybook_background(job_id: str):
                 job_id, 
                 JobStatus.PROCESSING, 
                 progress, 
-                f"Generating page {i+1}/{total_scenes}: {scene.visual[:50]}..."
+                f"Generating page {i+1}/{total_scenes}"
             )
             
             try:
-                print(f"🖼️ Generating page {i+1}")
-                print(f"📝 Pure prompt: {scene.visual}")
-                
-                # Generate image using pure prompt only
+                # Generate image
                 image = generate_image_simple(
                     scene.visual, 
                     story_request.model_choice, 
@@ -673,9 +635,8 @@ def generate_storybook_background(job_id: str):
                     story_request.consistency_seed
                 )
                 
-                # Save locally as backup
+                # Save locally
                 local_filepath, local_filename = save_image_to_local(image, scene.visual, story_request.style)
-                print(f"💾 Image saved locally as backup: {local_filename}")
                 
                 # Upload to HF Dataset
                 hf_url = None
@@ -690,161 +651,108 @@ def generate_storybook_background(job_id: str):
                     
                     if hf_url:
                         image_urls.append(hf_url)
-                        print(f"✅ Uploaded to HF Dataset: {hf_url}")
                 
-                # Store page data
                 page_data = {
                     "page_number": i + 1,
                     "image_url": hf_url or f"local://{local_filepath}",
-                    "hf_dataset_url": hf_url,
                     "text_content": scene.text,
-                    "visual_description": scene.visual,
-                    "prompt_used": scene.visual,
-                    "local_backup_path": local_filepath
+                    "visual_description": scene.visual
                 }
                 generated_pages.append(page_data)
                 
-                print(f"✅ Page {i+1} completed")
+                # Clean up after each page
+                if torch.cuda.is_available():
+                    torch.cuda.empty_cache()
+                gc.collect()
                 
             except Exception as e:
-                error_msg = f"Failed to generate page {i+1}: {str(e)}"
-                print(f"❌ {error_msg}")
-                update_job_status(job_id, JobStatus.FAILED, 0, error_msg)
+                print(f"❌ Page {i+1} failed: {e}")
+                update_job_status(job_id, JobStatus.FAILED, progress, str(e))
                 return
         
-        # Complete the job
         generation_time = time.time() - start_time
         
-        # Count successful HF uploads
-        hf_success_count = len(image_urls)
-        local_fallback_count = total_scenes - hf_success_count
-        
         result = {
             "story_title": story_request.story_title,
             "project_id": project_id,
             "total_pages": total_scenes,
-            "generated_pages": len(generated_pages),
             "generation_time": round(generation_time, 2),
             "hf_dataset_url": f"https://huggingface.co/datasets/{DATASET_ID}" if HF_TOKEN else None,
-            "consistency_seed": story_request.consistency_seed,
-            "pages": generated_pages,
             "image_urls": image_urls,
-            "upload_summary": {
-                "hf_successful": hf_success_count,
-                "local_fallback": local_fallback_count,
-                "total_attempted": total_scenes
-            }
+            "pages": generated_pages
         }
         
-        status_message = f"🎉 Storybook completed! {len(generated_pages)} pages created in {generation_time:.2f}s."
-        if hf_success_count > 0:
-            status_message += f" {hf_success_count} images uploaded to HF Dataset."
-        if local_fallback_count > 0:
-            status_message += f" {local_fallback_count} pages saved locally."
-        
         update_job_status(
             job_id, 
             JobStatus.COMPLETED, 
             100, 
-            status_message,
+            f"✅ Completed! {len(image_urls)} images uploaded",
             result
         )
         
-        print(f"🎉 Storybook generation finished for job {job_id}")
-        print(f"📤 HF Uploads: {hf_success_count} successful, {local_fallback_count} local fallbacks")
-        
     except Exception as e:
-        error_msg = f"Story generation failed: {str(e)}"
+        error_msg = f"Generation failed: {str(e)}"
         print(f"❌ {error_msg}")
         update_job_status(job_id, JobStatus.FAILED, 0, error_msg)
 
-# FASTAPI ENDPOINTS (for n8n)
+# FASTAPI ENDPOINTS
 @app.post("/api/generate-storybook")
 async def generate_storybook(request: dict, background_tasks: BackgroundTasks):
-    """Main endpoint for n8n integration - generates complete storybook using pure prompts"""
     try:
-        print(f"📥 Received n8n request for story: {request.get('story_title', 'Unknown')}")
+        print(f"📥 Received request for: {request.get('story_title', 'Unknown')}")
         
-        # Add consistency seed if not provided
-        if 'consistency_seed' not in request or not request['consistency_seed']:
+        if 'consistency_seed' not in request:
             request['consistency_seed'] = random.randint(1000, 9999)
-            print(f"🌱 Generated consistency seed: {request['consistency_seed']}")
         
-        # Generate project_id if not provided
         if 'project_id' not in request:
             request['project_id'] = request.get('story_title', 'unknown').replace(' ', '_').lower()
         
-        # Convert to Pydantic model
         story_request = StorybookRequest(**request)
         
-        # Validate required fields
         if not story_request.story_title or not story_request.scenes:
-            raise HTTPException(status_code=400, detail="story_title and scenes are required")
+            raise HTTPException(status_code=400, detail="story_title and scenes required")
         
-        # Create job immediately
         job_id = create_job(story_request)
-        
-        # Start background processing
         background_tasks.add_task(generate_storybook_background, job_id)
         
-        # Immediate response for n8n
-        response_data = {
+        return {
             "status": "success",
-            "message": "Storybook generation started",
             "job_id": job_id,
             "story_title": story_request.story_title,
             "project_id": request['project_id'],
             "total_scenes": len(story_request.scenes),
-            "consistency_seed": story_request.consistency_seed,
             "hf_dataset": f"https://huggingface.co/datasets/{DATASET_ID}" if HF_TOKEN else None,
-            "callback_url": story_request.callback_url,
-            "estimated_time_seconds": len(story_request.scenes) * 35,
-            "timestamp": datetime.now().isoformat()
+            "estimated_time_seconds": len(story_request.scenes) * 35
         }
         
-        print(f"✅ Job {job_id} started for: {story_request.story_title}")
-        
-        return response_data
-        
     except Exception as e:
-        error_msg = f"API Error: {str(e)}"
-        print(f"❌ {error_msg}")
-        raise HTTPException(status_code=500, detail=error_msg)
+        raise HTTPException(status_code=500, detail=str(e))
 
 @app.get("/api/job-status/{job_id}")
-async def get_job_status_endpoint(job_id: str):
-    """Check job status"""
+async def get_job_status(job_id: str):
     job_data = job_storage.get(job_id)
     if not job_data:
         raise HTTPException(status_code=404, detail="Job not found")
     
-    return JobStatusResponse(
-        job_id=job_id,
-        status=job_data["status"],
-        progress=job_data["progress"],
-        message=job_data["message"],
-        result=job_data["result"],
-        created_at=job_data["created_at"],
-        updated_at=job_data["updated_at"]
-    )
+    return {
+        "job_id": job_id,
+        "status": job_data["status"].value,
+        "progress": job_data["progress"],
+        "message": job_data["message"],
+        "result": job_data["result"]
+    }
 
 @app.get("/api/health")
-async def api_health():
-    """Health check endpoint for n8n"""
+async def health():
     return {
         "status": "healthy",
         "service": "storybook-generator",
         "hf_dataset": DATASET_ID if HF_TOKEN else "Disabled",
-        "hf_token_set": bool(HF_TOKEN),
-        "timestamp": datetime.now().isoformat(),
-        "active_jobs": len(job_storage),
-        "models_loaded": list(model_cache.keys())
+        "active_jobs": len(job_storage)
     }
 
-# NEW: Endpoint to get project images from HF Dataset
 @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"}
@@ -853,470 +761,65 @@ async def get_project_images(project_id: str):
         files = api.list_repo_files(repo_id=DATASET_ID, repo_type="dataset")
         
         project_files = [f for f in files if f.startswith(f"data/projects/{project_id}/")]
-        
         urls = [f"https://huggingface.co/datasets/{DATASET_ID}/resolve/main/{f}" for f in project_files]
         
-        return {
-            "project_id": project_id,
-            "total_images": len(urls),
-            "image_urls": urls
-        }
+        return {"project_id": project_id, "total_images": len(urls), "image_urls": urls}
     except Exception as e:
         return {"error": str(e)}
 
-# NEW MEMORY MANAGEMENT ENDPOINTS
-@app.get("/api/memory-status")
-async def get_memory_status():
-    """Get current memory usage and system status"""
-    memory_info = get_memory_usage()
-    return MemoryStatusResponse(
-        memory_used_mb=memory_info["memory_used_mb"],
-        memory_percent=memory_info["memory_percent"],
-        models_loaded=memory_info["models_loaded"],
-        active_jobs=memory_info["active_jobs"],
-        local_images_count=memory_info["local_images_count"],
-        gpu_memory_allocated_mb=memory_info["gpu_memory_allocated_mb"],
-        gpu_memory_cached_mb=memory_info["gpu_memory_cached_mb"],
-        status="healthy"
-    )
-
-@app.post("/api/clear-memory")
-async def clear_memory_endpoint(request: MemoryClearanceRequest):
-    """Clear memory by unloading models and cleaning up resources"""
-    try:
-        result = clear_memory(
-            clear_models=request.clear_models,
-            clear_jobs=request.clear_jobs,
-            clear_local_images=request.clear_local_images,
-            force_gc=request.force_gc
-        )
-        
-        return {
-            "status": "success",
-            "message": "Memory clearance completed",
-            "details": result
-        }
-        
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Memory clearance failed: {str(e)}")
-
-@app.post("/api/auto-cleanup")
-async def auto_cleanup():
-    """Automatic cleanup - clears completed jobs and forces GC"""
-    try:
-        result = clear_memory(
-            clear_models=False,  # Don't clear models by default
-            clear_jobs=True,     # Clear completed jobs
-            clear_local_images=False,  # Don't clear images by default
-            force_gc=True        # Force garbage collection
-        )
-        
-        return {
-            "status": "success",
-            "message": "Automatic cleanup completed",
-            "details": result
-        }
-        
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Auto cleanup failed: {str(e)}")
-
-@app.get("/api/local-images")
-async def get_local_images():
-    """API endpoint to get locally saved test images"""
-    storage_info = get_local_storage_info()
-    return storage_info
-
-@app.delete("/api/local-images/{filename:path}")
-async def delete_local_image_api(filename: str):
-    """API endpoint to delete a local image"""
-    try:
-        filepath = os.path.join(PERSISTENT_IMAGE_DIR, filename)
-        success, message = delete_local_image(filepath)
-        return {"status": "success" if success else "error", "message": message}
-    except Exception as e:
-        return {"status": "error", "message": str(e)}
-
-# SIMPLE GRADIO INTERFACE
+# GRADIO INTERFACE
 def create_gradio_interface():
-    """Create simple Gradio interface for testing"""
-    
-    def generate_test_image_simple(prompt, model_choice, style_choice):
-        """Generate a single image using pure prompt only"""
-        try:
-            if not prompt.strip():
-                return None, "❌ Please enter a prompt", None
-            
-            print(f"🎨 Generating test image with pure prompt: {prompt}")
-            
-            # Generate the image using pure prompt
-            image = generate_image_simple(
-                prompt, 
-                model_choice, 
-                style_choice,
-                1
-            )
-            
-            # Save to local storage
-            filepath, filename = save_image_to_local(image, prompt, style_choice)
-            
-            status_msg = f"""✅ Success! Generated: {prompt}
-
-📁 **Local file:** {filename if filename else 'Not saved'}"""
-            
-            return image, status_msg, filepath
-            
-        except Exception as e:
-            error_msg = f"❌ Generation failed: {str(e)}"
-            print(error_msg)
-            return None, error_msg, None
-    
-    with gr.Blocks(title="Simple Image Generator", theme="soft") as demo:
-        gr.Markdown("# 🎨 Simple Image Generator")
-        gr.Markdown("Generate images using **pure prompts only** - no automatic enhancements")
-        
-        # Storage info display
-        storage_info = gr.Textbox(
-            label="📊 Local Storage Information",
-            interactive=False,
-            lines=2
-        )
-        
-        # Memory status display
-        memory_status = gr.Textbox(
-            label="🧠 Memory Status",
-            interactive=False,
-            lines=3
-        )
-        
-        # HF Dataset status
-        hf_status = gr.Textbox(
-            label="📤 Hugging Face Dataset",
-            value=f"✅ Connected to {DATASET_ID}" if HF_TOKEN else "❌ HF_TOKEN not set - local only",
-            interactive=False,
-            lines=2
-        )
+    def generate_test(prompt, model_choice, style_choice):
+        if not prompt.strip():
+            return None, "❌ Please enter a prompt"
         
-        def update_storage_info():
-            info = get_local_storage_info()
-            if "error" not in info:
-                return f"📁 Local Storage: {info['total_files']} images, {info['total_size_mb']} MB used"
-            return "📁 Local Storage: Unable to calculate"
+        image = generate_image_simple(prompt, model_choice, style_choice, 1)
+        filepath, filename = save_image_to_local(image, prompt, style_choice)
         
-        def update_memory_status():
-            memory_info = get_memory_usage()
-            status_text = f"🧠 Memory Usage: {memory_info['memory_used_mb']} MB ({memory_info['memory_percent']}%)\n"
-            status_text += f"📦 Models Loaded: {memory_info['models_loaded']}\n"
-            status_text += f"⚡ Active Jobs: {memory_info['active_jobs']}"
-            
-            if memory_info['gpu_memory_allocated_mb']:
-                status_text += f"\n🎮 GPU Memory: {memory_info['gpu_memory_allocated_mb']} MB allocated"
-            
-            return status_text
+        return image, f"✅ Generated! Local: {filename}"
+    
+    with gr.Blocks(title="Storybook Generator") as demo:
+        gr.Markdown("# 🎨 Storybook Generator")
         
         with gr.Row():
-            with gr.Column(scale=1):
-                gr.Markdown("### 🎯 Quality Settings")
-                
-                model_dropdown = gr.Dropdown(
-                    label="AI Model",
-                    choices=list(MODEL_CHOICES.keys()),
-                    value="dreamshaper-8"
-                )
-                
-                style_dropdown = gr.Dropdown(
-                    label="Art Style",
-                    choices=["childrens_book", "realistic", "fantasy", "anime"],
-                    value="anime"
-                )
-                
-                prompt_input = gr.Textbox(
-                    label="Pure Prompt",
-                    placeholder="Enter your exact prompt...",
-                    lines=3
-                )
-                
-                generate_btn = gr.Button("✨ Generate Image", variant="primary")
-                
-                # Current image management
-                current_file_path = gr.State()
-                delete_btn = gr.Button("🗑️ Delete This Image", variant="stop")
-                delete_status = gr.Textbox(label="Delete Status", interactive=False, lines=2)
-                
-                # Memory management section
-                gr.Markdown("### 🧠 Memory Management")
-                with gr.Row():
-                    auto_cleanup_btn = gr.Button("🔄 Auto Cleanup", size="sm")
-                    clear_models_btn = gr.Button("🗑️ Clear Models", variant="stop", size="sm")
-                
-                memory_clear_status = gr.Textbox(label="Memory Clear Status", interactive=False, lines=2)
-                
-                gr.Markdown("### 📚 API Usage for n8n")
-                gr.Markdown(f"""
-                **Generate Storybook:**
-                - Endpoint: `POST /api/generate-storybook`
-                - Body: `{{"story_title": "...", "scenes": [...]}}`
-                
-                **Check Status:**
-                - `GET /api/job-status/{{job_id}}`
-                
-                **HF Dataset:**
-                - `{DATASET_ID if HF_TOKEN else "Set HF_TOKEN to enable"}`
-                """)
-            
-            with gr.Column(scale=2):
-                image_output = gr.Image(label="Generated Image", height=500, show_download_button=True)
-                status_output = gr.Textbox(label="Status", interactive=False, lines=4)
-        
-        # Local file management section
-        with gr.Accordion("📁 Manage Local Test Images", open=True):
-            gr.Markdown("### Locally Saved Images")
-            
-            with gr.Row():
-                refresh_btn = gr.Button("🔄 Refresh List")
-                clear_all_btn = gr.Button("🗑️ Clear All Images", variant="stop")
+            with gr.Column():
+                model = gr.Dropdown(choices=list(MODEL_CHOICES.keys()), value="dreamshaper-8", label="Model")
+                style = gr.Dropdown(choices=["childrens_book", "realistic", "fantasy", "anime"], value="anime", label="Style")
+                prompt = gr.Textbox(label="Prompt", lines=3)
+                btn = gr.Button("Generate", variant="primary")
             
-            file_gallery = gr.Gallery(
-                label="Local Images",
-                show_label=True,
-                elem_id="gallery",
-                columns=4,
-                height="auto"
-            )
-            
-            clear_status = gr.Textbox(label="Clear Status", interactive=False)
-        
-        def delete_current_image(filepath):
-            """Delete the currently displayed image"""
-            if not filepath:
-                return "❌ No image to delete", None, None, refresh_local_images()
-            
-            success, message = delete_local_image(filepath)
-            updated_files = refresh_local_images()
-            
-            if success:
-                status_msg = f"✅ {message}"
-                return status_msg, None, "Image deleted successfully!", updated_files
-            else:
-                return f"❌ {message}", None, "Delete failed", updated_files
-
-        def clear_all_images():
-            """Delete all local images"""
-            try:
-                storage_info = get_local_storage_info()
-                deleted_count = 0
-                
-                if "images" in storage_info:
-                    for image_info in storage_info["images"]:
-                        success, _ = delete_local_image(image_info["path"])
-                        if success:
-                            deleted_count += 1
-                
-                updated_files = refresh_local_images()
-                return f"✅ Deleted {deleted_count} images", updated_files
-            except Exception as e:
-                return f"❌ Error: {str(e)}", refresh_local_images()
-        
-        def perform_auto_cleanup():
-            """Perform automatic cleanup"""
-            try:
-                result = clear_memory(
-                    clear_models=False,
-                    clear_jobs=True,
-                    clear_local_images=False,
-                    force_gc=True
-                )
-                return f"✅ Auto cleanup completed: {len(result['actions_performed'])} actions"
-            except Exception as e:
-                return f"❌ Auto cleanup failed: {str(e)}"
-        
-        def clear_models():
-            """Clear all loaded models"""
-            try:
-                result = clear_memory(
-                    clear_models=True,
-                    clear_jobs=False,
-                    clear_local_images=False,
-                    force_gc=True
-                )
-                return f"✅ Models cleared: {len(result['actions_performed'])} actions"
-            except Exception as e:
-                return f"❌ Model clearance failed: {str(e)}"
-
-        # Connect buttons to functions
-        generate_btn.click(
-            fn=generate_test_image_simple,
-            inputs=[prompt_input, model_dropdown, style_dropdown],
-            outputs=[image_output, status_output, current_file_path]
-        ).then(
-            fn=refresh_local_images,
-            outputs=file_gallery
-        ).then(
-            fn=update_storage_info,
-            outputs=storage_info
-        ).then(
-            fn=update_memory_status,
-            outputs=memory_status
-        )
-        
-        delete_btn.click(
-            fn=delete_current_image,
-            inputs=current_file_path,
-            outputs=[delete_status, image_output, status_output, file_gallery]
-        ).then(
-            fn=update_storage_info,
-            outputs=storage_info
-        ).then(
-            fn=update_memory_status,
-            outputs=memory_status
-        )
-        
-        refresh_btn.click(
-            fn=refresh_local_images,
-            outputs=file_gallery
-        ).then(
-            fn=update_storage_info,
-            outputs=storage_info
-        ).then(
-            fn=update_memory_status,
-            outputs=memory_status
-        )
+            with gr.Column():
+                output = gr.Image(label="Generated Image", height=500)
+                status = gr.Textbox(label="Status")
         
-        clear_all_btn.click(
-            fn=clear_all_images,
-            outputs=[clear_status, file_gallery]
-        ).then(
-            fn=update_storage_info,
-            outputs=storage_info
-        ).then(
-            fn=update_memory_status,
-            outputs=memory_status
-        )
-        
-        # Memory management buttons
-        auto_cleanup_btn.click(
-            fn=perform_auto_cleanup,
-            outputs=memory_clear_status
-        ).then(
-            fn=update_memory_status,
-            outputs=memory_status
-        )
-        
-        clear_models_btn.click(
-            fn=clear_models,
-            outputs=memory_clear_status
-        ).then(
-            fn=update_memory_status,
-            outputs=memory_status
-        )
-        
-        # Initialize on load
-        demo.load(fn=refresh_local_images, outputs=file_gallery)
-        demo.load(fn=update_storage_info, outputs=storage_info)
-        demo.load(fn=update_memory_status, outputs=memory_status)
+        btn.click(fn=generate_test, inputs=[prompt, model, style], outputs=[output, status])
     
     return demo
 
-# Create simple Gradio app
 demo = create_gradio_interface()
 
-# Simple root endpoint
 @app.get("/")
 async def root():
     return {
-        "message": "Storybook Generator API with HF Dataset is running!",
-        "api_endpoints": {
-            "health_check": "GET /api/health",
-            "generate_storybook": "POST /api/generate-storybook",
-            "check_job_status": "GET /api/job-status/{job_id}",
-            "project_images": "GET /api/project-images/{project_id}",
-            "local_images": "GET /api/local-images",
-            "memory_status": "GET /api/memory-status",
-            "clear_memory": "POST /api/clear-memory",
-            "auto_cleanup": "POST /api/auto-cleanup"
-        },
+        "message": "Storybook Generator API",
         "hf_dataset": DATASET_ID if HF_TOKEN else "Disabled",
-        "features": {
-            "pure_prompts": "✅ Enabled - No automatic enhancements",
-            "n8n_integration": "✅ Enabled",
-            "memory_management": "✅ Enabled",
-            "hf_dataset": "✅ Enabled" 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}"
         },
-        "web_interface": "GET /ui"
-    }
-
-# Add a simple test endpoint
-@app.get("/api/test")
-async def test_endpoint():
-    return {
-        "status": "success", 
-        "message": "API is working correctly",
-        "hf_dataset": DATASET_ID if HF_TOKEN else "Disabled",
-        "timestamp": datetime.now().isoformat()
+        "ui": "/ui"
     }
 
-# For Hugging Face Spaces deployment
-def get_app():
-    return app
-
 if __name__ == "__main__":
     import uvicorn
-    import os
-    
-    # Check if we're running on Hugging Face Spaces
-    HF_SPACE = os.environ.get('SPACE_ID') is not None
     
-    if HF_SPACE:
-        print("🚀 Running on Hugging Face Spaces - Integrated Mode")
+    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 available at: /api/*")
-        print("🎨 Web interface available at: /ui")
-        print("📝 PURE PROMPTS enabled - no automatic enhancements")
-        print("🧠 MEMORY MANAGEMENT enabled - automatic cleanup available")
-        
-        # Mount Gradio without reassigning app
         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"
-        )
+        uvicorn.run(app, host="0.0.0.0", port=7860, log_level="info")
     else:
-        # Local development - run separate servers
-        print("🚀 Running locally - Separate API and UI servers")
-        print(f"📦 HF Dataset: {DATASET_ID if HF_TOKEN else 'Disabled'}")
-        print("📚 API endpoints: http://localhost:8000/api/*")
-        print("🎨 Web interface: http://localhost:7860/ui")
-        print("📝 PURE PROMPTS enabled - no automatic enhancements")
-        print("🧠 MEMORY MANAGEMENT enabled - automatic cleanup available")
-        
-        def run_fastapi():
-            """Run FastAPI on port 8000 for API calls"""
-            uvicorn.run(
-                app, 
-                host="0.0.0.0", 
-                port=8000, 
-                log_level="info",
-                access_log=False
-            )
-        
-        def run_gradio():
-            """Run Gradio on port 7860 for web interface"""
-            demo.launch(server_name="0.0.0.0", server_port=7860, share=False)
-        
-        # Run both servers in separate threads
-        import threading
-        fastapi_thread = threading.Thread(target=run_fastapi, daemon=True)
-        gradio_thread = threading.Thread(target=run_gradio, daemon=True)
-        
-        fastapi_thread.start()
-        gradio_thread.start()
-        
-        try:
-            # Keep main thread alive
-            while True:
-                time.sleep(1)
-        except KeyboardInterrupt:
-            print("🛑 Shutting down servers...")
+        print("🚀 Running locally")
+        uvicorn.run(app, host="0.0.0.0", port=8000)