Update app.py

app.py

@@ -10,11 +10,14 @@ import re
 import time
 import json
 from typing import List, Optional, Dict
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI, HTTPException, BackgroundTasks
 from pydantic import BaseModel
 import gc
 import psutil
 import threading
+import uuid
+import hashlib
+from enum import Enum
 
 # External OCI API URL
 OCI_API_BASE_URL = "https://yukee1992-oci-story-book.hf.space"
@@ -32,6 +35,13 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
+# Job Status Enum
+class JobStatus(str, Enum):
+    PENDING = "pending"
+    PROCESSING = "processing"
+    COMPLETED = "completed"
+    FAILED = "failed"
+
 # Story scene model
 class StoryScene(BaseModel):
     visual: str
@@ -47,6 +57,7 @@ class StorybookRequest(BaseModel):
     characters: List[CharacterDescription] = []
     model_choice: str = "dreamshaper-8"
     style: str = "childrens_book"
+    callback_url: Optional[str] = None  # For n8n webhook
 
 class StorybookResponse(BaseModel):
     status: str
@@ -60,22 +71,35 @@ class StorybookResponse(BaseModel):
     pages: List[dict]
     request_id: str
 
-# MODEL SELECTION
+class JobStatusResponse(BaseModel):
+    job_id: str
+    status: JobStatus
+    progress: int  # 0-100
+    message: str
+    result: Optional[dict] = None
+    created_at: float
+    updated_at: float
+
+# HIGH-QUALITY MODEL SELECTION (from your high-quality app.py)
 MODEL_CHOICES = {
-    "dreamshaper-8": "lykon/dreamshaper-8",
-    "realistic-vision": "SG161222/Realistic_Vision_V5.1", 
+    "dreamshaper-8": "lykon/dreamshaper-8",  # Excellent for fantasy/artistic
+    "realistic-vision": "SG161222/Realistic_Vision_V5.1",  # Photorealistic
+    "anything-v5": "andite/anything-v5.0",  # Anime/illustration style
+    "openjourney": "prompthero/openjourney",  # Artistic/Midjourney style
+    "sd-2.1": "stabilityai/stable-diffusion-2-1",  # General purpose
 }
 
+# GLOBAL STORAGE FOR JOBS AND STATUS
+job_storage = {}
+active_requests = {}
+character_seeds = {}
+
 # GLOBAL MODEL CACHE
 model_cache = {}
 current_model_name = None
 current_pipe = None
 model_lock = threading.Lock()
 
-# Character consistency tracking
-character_seeds = {}
-active_requests = {}
-
 def monitor_memory():
     try:
         process = psutil.Process()
@@ -89,7 +113,7 @@ def cleanup_memory():
         torch.cuda.empty_cache()
 
 def load_model(model_name="dreamshaper-8"):
-    """Thread-safe model loading"""
+    """Thread-safe model loading with HIGH-QUALITY settings"""
     global model_cache, current_model_name, current_pipe
     
     with model_lock:
@@ -98,10 +122,11 @@ def load_model(model_name="dreamshaper-8"):
             current_model_name = model_name
             return current_pipe
         
-        print(f"🔄 Loading model: {model_name}")
+        print(f"🔄 Loading HIGH-QUALITY model: {model_name}")
         try:
             model_id = MODEL_CHOICES.get(model_name, "lykon/dreamshaper-8")
             
+            # HIGH-QUALITY LOADING (from your high-quality app.py)
             pipe = StableDiffusionPipeline.from_pretrained(
                 model_id, 
                 torch_dtype=torch.float32,
@@ -109,6 +134,7 @@ def load_model(model_name="dreamshaper-8"):
                 requires_safety_checker=False
             )
             
+            # Use better scheduler for quality
             pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
             pipe = pipe.to("cpu")
             
@@ -116,352 +142,457 @@ def load_model(model_name="dreamshaper-8"):
             current_pipe = pipe
             current_model_name = model_name
             
-            print(f"✅ Model loaded: {model_name}")
+            print(f"✅ HIGH-QUALITY Model loaded: {model_name}")
             return pipe
             
         except Exception as e:
             print(f"❌ Model loading failed: {e}")
-            raise
+            # Fallback to SD 1.5
+            return StableDiffusionPipeline.from_pretrained(
+                "runwayml/stable-diffusion-v1-5", 
+                torch_dtype=torch.float32,
+                safety_checker=None,
+                requires_safety_checker=False
+            ).to("cpu")
 
 # Initialize default model
-print("🚀 Initializing Storybook Generator...")
+print("🚀 Initializing HIGH-QUALITY Storybook Generator...")
 load_model("dreamshaper-8")
-print("✅ Model loaded and ready!")
+print("✅ HIGH-QUALITY Model loaded and ready!")
+
+# PROFESSIONAL PROMPT ENGINEERING (from your high-quality app.py)
+def enhance_prompt(prompt, style="childrens_book"):
+    """Transform basic prompts into professional-grade prompts"""
+    
+    style_templates = {
+        "childrens_book": [
+            "masterpiece, best quality, 4K, ultra detailed, children's book illustration",
+            "watercolor painting, whimsical, cute, charming, storybook style",
+            "vibrant colors, soft lighting, magical, enchanting, dreamlike",
+            "Pixar style, Disney animation, high detail, professional artwork"
+        ],
+        "realistic": [
+            "photorealistic, 8K, ultra detailed, professional photography",
+            "sharp focus, studio lighting, high resolution, intricate details",
+            "realistic textures, natural lighting, cinematic quality"
+        ],
+        "fantasy": [
+            "epic fantasy art, digital painting, concept art, trending on artstation",
+            "magical, mystical, ethereal, otherworldly, fantasy illustration",
+            "dynamic composition, dramatic lighting, highly detailed"
+        ],
+        "anime": [
+            "anime style, Japanese animation, high quality, detailed artwork",
+            "beautiful anime illustration, vibrant colors, clean lines",
+            "studio ghibli style, makoto shinkai, professional anime art"
+        ]
+    }
+    
+    # Choose style template
+    templates = style_templates.get(style, style_templates["childrens_book"])
+    style_prompt = templates[0]  # Use the first template
+    
+    # Enhanced prompt construction
+    enhanced = f"{style_prompt}, {prompt}"
+    
+    # Add quality boosters
+    quality_boosters = [
+        "intricate details", "beautiful composition", "perfect lighting",
+        "professional artwork", "award winning", "trending on artstation"
+    ]
+    
+    # Add 2-3 random quality boosters
+    import random
+    boosters = random.sample(quality_boosters, 2)
+    enhanced += ", " + ", ".join(boosters)
+    
+    # Negative prompt to avoid bad quality
+    negative_prompt = (
+        "blurry, low quality, low resolution, ugly, deformed, poorly drawn, "
+        "bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, "
+        "disconnected limbs, mutation, mutated, ugly, disgusting, bad art, "
+        "beginner, amateur, distorted, watermark, signature, text, username"
+    )
+    
+    return enhanced, negative_prompt
 
-# PROMPT OPTIMIZATION - PRESERVE FULL DESCRIPTIONS
 def optimize_prompt(scene_visual, characters, style="childrens_book", page_number=1):
     """
-    Create a prompt that PRESERVES all visual descriptions while fitting 77 tokens
+    Create HIGH-QUALITY prompt that preserves visual descriptions
     """
-    # 1. PRESERVE THE ENTIRE SCENE VISUAL DESCRIPTION (most important)
-    scene_prompt = scene_visual
+    # First, enhance the base prompt using professional engineering
+    base_enhanced, _ = enhance_prompt(scene_visual, style)
     
-    # 2. Extract only ESSENTIAL character features (not full descriptions)
+    # Extract character essence
     character_essence = ""
     if characters:
         char_names = []
         for char in characters:
             char_name = char.get('name', '') if isinstance(char, dict) else getattr(char, 'name', '')
-            char_names.append(char_name.split()[0])  # Just first name
+            char_names.append(char_name.split()[0])
         
         character_essence = f" featuring {', '.join(char_names)}"
     
-    # 3. Add style context briefly
-    style_context = {
-        "childrens_book": "children's book illustration style",
-        "realistic": "photorealistic style",
-        "fantasy": "fantasy art style",
-        "anime": "anime style"
-    }.get(style, "children's book illustration style")
-    
-    # 4. Build the final prompt - SCENE DESCRIPTION COMES FIRST
+    # Build final HIGH-QUALITY prompt
     continuity = f"Scene {page_number}, " if page_number > 1 else ""
-    final_prompt = f"{continuity}{scene_prompt}{character_essence}. {style_context}. high quality, detailed"
-    
-    # 5. If still too long, prioritize scene description over style
-    words = final_prompt.split()
-    if len(words) > 60:
-        # Keep the scene description intact, trim the end
-        scene_words = scene_visual.split()
-        if len(scene_words) > 45:
-            # If scene itself is too long, keep first 40 words of scene
-            scene_part = ' '.join(scene_words[:40])
-            final_prompt = f"{continuity}{scene_part}...{character_essence}. {style_context}"
-        else:
-            # Keep entire scene, trim style part
-            final_prompt = f"{continuity}{scene_visual}{character_essence}. high quality"
-    
-    print(f"📝 Final prompt: {final_prompt}")
-    print(f"📏 Length: {len(final_prompt.split())} words")
+    final_prompt = f"{continuity}{base_enhanced}{character_essence}"
     
+    print(f"🎨 HIGH-QUALITY Prompt: {final_prompt}")
     return final_prompt
 
-def enhance_prompt(scene_visual, characters, style="childrens_book", page_number=1):
-    """Create optimized prompt that preserves visual descriptions"""
-    main_prompt = optimize_prompt(scene_visual, characters, style, page_number)
-    
-    negative_prompt = (
-        "blurry, low quality, ugly, deformed, bad anatomy, "
-        "watermark, text, username, multiple people, inconsistent, "
-        "missing limbs, extra limbs, disfigured, malformed"
-    )
-    
-    return main_prompt, negative_prompt
+def generate_high_quality_image(prompt, model_choice="dreamshaper-8", style="childrens_book", negative_prompt=""):
+    """Generate HIGH-QUALITY image with professional settings"""
+    try:
+        # Load the appropriate model
+        pipe = load_model(model_choice)
+        
+        # Generate HIGH-QUALITY image with better settings
+        image = pipe(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            num_inference_steps=30,  # More steps = better quality
+            guidance_scale=8.5,      # Higher guidance = better prompt following
+            width=768,               # Higher resolution
+            height=768,
+            generator=torch.Generator(device="cpu").manual_seed(int(time.time()))
+        ).images[0]
+        
+        print("✅ HIGH-QUALITY Image generated successfully!")
+        return image
+        
+    except Exception as e:
+        print(f"❌ HQ Generation failed: {str(e)}")
+        raise
 
-def save_complete_storybook_page(image, story_title, sequence_number, scene_text):
+def save_to_oci_via_api(image, prompt, story_title, page_number):
+    """Save image using the EXTERNAL OCI API endpoint"""
     try:
+        # Convert image to bytes
         img_bytes = io.BytesIO()
         image.save(img_bytes, format='PNG')
         img_data = img_bytes.getvalue()
         
-        clean_title = re.sub(r'[^a-zA-Z0-9_\-]', '', story_title.strip().replace(' ', '_'))
-        image_filename = f"page_{sequence_number:03d}_{clean_title}.png"
-        text_filename = f"page_{sequence_number:03d}_{clean_title}.txt"
+        # Create filename
+        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"page_{page_number:03d}_{safe_prompt}_{timestamp}.png"
+        
+        # Use the EXTERNAL API URL
+        api_url = f"{OCI_API_BASE_URL}/api/upload"
         
-        # Save image to OCI
-        image_api_url = f"{OCI_API_BASE_URL}/api/upload"
-        files = {'file': (image_filename, img_data, 'image/png')}
-        data = {'project_id': 'storybook-library', 'subfolder': f'stories/{clean_title}'}
-        response = requests.post(image_api_url, files=files, data=data, timeout=30)
+        print(f"🌐 Calling EXTERNAL OCI API: {api_url}")
+        
+        # Prepare form data for API request
+        files = {
+            'file': (filename, img_data, 'image/png')
+        }
         
-        # Save text to OCI
-        text_api_url = f"{OCI_API_BASE_URL}/api/upload-text"
-        text_data = {
+        data = {
             'project_id': 'storybook-library',
-            'subfolder': f'stories/{clean_title}',
-            'filename': text_filename,
-            'content': scene_text
+            'subfolder': f'stories/{story_title}'
         }
-        text_response = requests.post(text_api_url, data=text_data, timeout=30)
         
-        return True, f"✅ Page {sequence_number}: Image & Text saved"
+        # Make the API request with timeout
+        response = requests.post(api_url, files=files, data=data, timeout=30)
+        
+        print(f"📨 API Response: {response.status_code}")
+        
+        if response.status_code == 200:
+            result = response.json()
+            if result['status'] == 'success':
+                return result.get('file_url', 'Unknown URL')
+            else:
+                raise Exception(f"API Error: {result.get('message', 'Unknown error')}")
+        else:
+            raise Exception(f"HTTP Error: {response.status_code}")
             
     except Exception as e:
-        return False, f"❌ Save failed: {str(e)}"
+        raise Exception(f"API upload failed: {str(e)}")
+
+# JOB MANAGEMENT FUNCTIONS
+def create_job(story_request: StorybookRequest) -> str:
+    """Create a new job and return job ID"""
+    job_id = str(uuid.uuid4())
+    
+    job_storage[job_id] = {
+        "status": JobStatus.PENDING,
+        "progress": 0,
+        "message": "Job created and queued",
+        "request": story_request.dict(),
+        "result": None,
+        "created_at": time.time(),
+        "updated_at": time.time(),
+        "pages": []
+    }
+    
+    print(f"📝 Created job {job_id} for story: {story_request.story_title}")
+    return job_id
+
+def update_job_status(job_id: str, status: JobStatus, progress: int, message: str, result=None):
+    """Update job status and send notification if callback URL exists"""
+    if job_id not in job_storage:
+        return False
+    
+    job_storage[job_id].update({
+        "status": status,
+        "progress": progress,
+        "message": message,
+        "updated_at": time.time()
+    })
+    
+    if result:
+        job_storage[job_id]["result"] = result
+        if "pages" in result:
+            job_storage[job_id]["pages"] = result["pages"]
+    
+    # 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_data = {
+                "job_id": job_id,
+                "status": status.value,
+                "progress": progress,
+                "message": message,
+                "story_title": request_data["story_title"],
+                "timestamp": time.time()
+            }
+            
+            if status == JobStatus.COMPLETED and result:
+                callback_data["result"] = result
+            
+            response = requests.post(
+                request_data["callback_url"],
+                json=callback_data,
+                timeout=10
+            )
+            print(f"📢 Callback sent to {request_data['callback_url']}: Status {response.status_code}")
+        except Exception as e:
+            print(f"⚠️ Callback failed: {str(e)}")
+    
+    return True
 
-def get_character_seed(story_title, character_name, page_number):
-    if story_title not in character_seeds:
-        character_seeds[story_title] = {}
+def get_job_status(job_id: str) -> Optional[JobStatusResponse]:
+    """Get current job status"""
+    if job_id not in job_storage:
+        return None
     
-    seed_key = f"{character_name}_{page_number}"
-    if seed_key not in character_seeds[story_title]:
-        base_seed = hash(f"{story_title}_{character_name}") % 1000000
-        page_variation = (page_number * 13) % 1000
-        seed_value = (base_seed + page_variation) % 1000000
-        character_seeds[story_title][seed_key] = seed_value
+    job_data = job_storage[job_id]
     
-    return character_seeds[story_title][seed_key]
+    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"]
+    )
 
-def process_storybook_generation(request_id, request_data):
-    """Process generation in background and store results"""
+# BACKGROUND TASK FOR STORY GENERATION
+def generate_storybook_background(job_id: str):
+    """Background task to generate storybook pages"""
     try:
-        print(f"🔧 Processing request {request_id} in background...")
-        
-        # Load model
-        load_model(request_data["model_choice"])
+        job_data = job_storage[job_id]
+        story_request_data = job_data["request"]
+        story_request = StorybookRequest(**story_request_data)
         
-        # Convert characters to dict
-        characters_dict = []
-        for char in request_data["characters"]:
-            characters_dict.append({
-                "name": char["name"],
-                "description": char["description"]
-            })
+        print(f"🎬 Starting story generation for job {job_id}")
+        update_job_status(job_id, JobStatus.PROCESSING, 10, "Starting storybook generation...")
         
-        results = []
-        status_messages = []
+        total_scenes = len(story_request.scenes)
+        generated_pages = []
         start_time = time.time()
         
-        # Process each page
-        for i, scene in enumerate(request_data["scenes"], 1):
+        for i, scene in enumerate(story_request.scenes):
+            progress = 10 + int((i / total_scenes) * 80)
+            update_job_status(
+                job_id, 
+                JobStatus.PROCESSING, 
+                progress, 
+                f"Generating page {i+1}/{total_scenes}: {scene.visual[:50]}..."
+            )
+            
             try:
-                print(f"🔄 Generating page {i}...")
+                # Create HIGH-QUALITY prompt
+                prompt, negative_prompt = enhance_prompt(scene.visual, story_request.style)
+                prompt = optimize_prompt(scene.visual, story_request.characters, story_request.style, i+1)
                 
-                enhanced_prompt, negative_prompt = enhance_prompt(
-                    scene["visual"], characters_dict, request_data["style"], i
+                # Generate HIGH-QUALITY image
+                image = generate_high_quality_image(
+                    prompt, 
+                    story_request.model_choice, 
+                    story_request.style,
+                    negative_prompt
                 )
                 
-                # Get character name for seed
-                main_char_name = "default"
-                if characters_dict:
-                    main_char_name = characters_dict[0]["name"]
-                
-                # Use consistent seed
-                generator = torch.Generator(device="cpu")
-                main_char_seed = get_character_seed(request_data["story_title"], main_char_name, i)
-                generator.manual_seed(main_char_seed)
-                
-                # Generate image
-                global current_pipe
-                image = current_pipe(
-                    prompt=enhanced_prompt,
-                    negative_prompt=negative_prompt,
-                    num_inference_steps=25,
-                    guidance_scale=7.0,
-                    width=512,
-                    height=512,
-                    generator=generator
-                ).images[0]
-                
                 # Save to OCI
-                success, save_status = save_complete_storybook_page(
-                    image, request_data["story_title"], i, scene["text"]
-                )
+                image_url = save_to_oci_via_api(image, prompt, story_request.story_title, i+1)
                 
-                if success:
-                    results.append({"page_number": i, "status": "success"})
-                    status_messages.append(f"Page {i}: {save_status}")
-                    print(f"✅ Page {i} completed")
-                else:
-                    results.append({"page_number": i, "status": "error", "message": save_status})
-                    status_messages.append(f"Page {i}: {save_status}")
+                # Store page data
+                page_data = {
+                    "page_number": i + 1,
+                    "image_url": image_url,
+                    "text": scene.text,
+                    "prompt": prompt
+                }
+                generated_pages.append(page_data)
                 
-                cleanup_memory()
+                print(f"✅ Page {i+1} completed successfully")
                 
-                if i < len(request_data["scenes"]):
-                    time.sleep(1)
-                    
             except Exception as e:
-                error_msg = f"Page {i} failed: {str(e)}"
-                results.append({"page_number": i, "status": "error", "message": error_msg})
-                status_messages.append(error_msg)
+                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)
+                return
         
-        total_time = time.time() - start_time
-        
-        # Store results
-        active_requests[request_id] = {
-            "status": "completed",
-            "results": results,
-            "message": "\n".join(status_messages),
-            "generation_time": total_time,
-            "completed_at": datetime.now().isoformat()
+        # Complete the job
+        generation_time = time.time() - start_time
+        result = {
+            "story_title": story_request.story_title,
+            "total_pages": total_scenes,
+            "characters_used": len(story_request.characters),
+            "generated_pages": len(generated_pages),
+            "generation_time": round(generation_time, 2),
+            "folder_path": f"stories/{story_request.story_title}",
+            "pages": generated_pages
         }
         
-        print(f"✅ Request {request_id} completed in {total_time:.2f} seconds")
+        update_job_status(
+            job_id, 
+            JobStatus.COMPLETED, 
+            100, 
+            f"Storybook generation completed! {len(generated_pages)} pages created in {generation_time:.2f}s",
+            result
+        )
+        
+        print(f"🎉 Storybook generation completed for job {job_id}")
         
     except Exception as e:
-        active_requests[request_id] = {
-            "status": "error",
-            "message": f"Processing failed: {str(e)}"
-        }
-        print(f"❌ Request {request_id} failed: {e}")
+        error_msg = f"Story generation failed: {str(e)}"
+        print(f"❌ {error_msg}")
+        update_job_status(job_id, JobStatus.FAILED, 0, error_msg)
 
-# FastAPI endpoint - IMMEDIATE RESPONSE
-@app.post("/api/generate-storybook", response_model=StorybookResponse)
-async def api_generate_storybook(request: StorybookRequest):
-    """API endpoint that returns immediately"""
+# FASTAPI ENDPOINTS
+@app.post("/api/generate-storybook", response_model=dict)
+async def generate_storybook(request: StorybookRequest, background_tasks: BackgroundTasks):
+    """Endpoint to start storybook generation - immediately returns job ID"""
     try:
-        print(f"📚 Received request: {request.story_title}")
-        print(f"📖 Pages: {len(request.scenes)}")
-        
-        # Create request ID
-        request_id = f"{request.story_title}_{int(time.time())}"
-        
-        # Convert to dict for background processing
-        request_data = {
-            "story_title": request.story_title,
-            "scenes": [{"visual": scene.visual, "text": scene.text} for scene in request.scenes],
-            "characters": [{"name": char.name, "description": char.description} for char in request.characters],
-            "model_choice": request.model_choice,
-            "style": request.style
-        }
-        
-        # Store initial request state
-        active_requests[request_id] = {
-            "status": "processing",
-            "started_at": datetime.now().isoformat(),
-            "total_pages": len(request.scenes)
-        }
+        # Create job immediately
+        job_id = create_job(request)
         
-        # Start background processing in a thread
-        import threading
-        thread = threading.Thread(
-            target=process_storybook_generation,
-            args=(request_id, request_data)
-        )
-        thread.daemon = True
-        thread.start()
+        # Start background task
+        background_tasks.add_task(generate_storybook_background, job_id)
         
-        # IMMEDIATE RESPONSE to n8n
-        response_data = {
-            "status": "processing",
+        # Return job ID immediately for status tracking
+        return {
+            "status": "success",
+            "message": "Storybook generation started",
+            "job_id": job_id,
             "story_title": request.story_title,
-            "total_pages": len(request.scenes),
-            "characters_used": len(request.characters),
-            "generated_pages": 0,
-            "generation_time": 0,
-            "message": f"Generation started for {len(request.scenes)} pages. Request ID: {request_id}",
-            "folder_path": f"storybook-library/stories/{request.story_title.replace(' ', '_')}/",
-            "pages": [
-                {
-                    "page_number": i+1,
-                    "image_file": f"page_{i+1:03d}_{request.story_title.replace(' ', '_')}.png",
-                    "text_file": f"page_{i+1:03d}_{request.story_title.replace(' ', '_')}.txt"
-                } for i in range(len(request.scenes))
-            ],
-            "request_id": request_id
+            "total_scenes": len(request.scenes),
+            "callback_url": request.callback_url
         }
         
-        return response_data
-        
     except Exception as e:
-        error_msg = f"Request failed: {str(e)}"
-        print(f"❌ {error_msg}")
-        raise HTTPException(status_code=500, detail=error_msg)
+        raise HTTPException(status_code=500, detail=f"Failed to start generation: {str(e)}")
 
-# Status check endpoint for n8n
-@app.get("/api/status/{request_id}")
-async def check_status(request_id: str):
-    """Check status of a generation request"""
-    if request_id not in active_requests:
-        return {"status": "not_found", "message": "Request ID not found"}
-    
-    request_data = active_requests[request_id]
-    return {
-        "status": request_data["status"],
-        "message": request_data.get("message", ""),
-        "generation_time": request_data.get("generation_time", 0),
-        "completed_at": request_data.get("completed_at", ""),
-        "total_pages": request_data.get("total_pages", 0)
-    }
+@app.get("/api/job-status/{job_id}", response_model=JobStatusResponse)
+async def get_job_status_endpoint(job_id: str):
+    """Endpoint to check job status"""
+    job_status = get_job_status(job_id)
+    if not job_status:
+        raise HTTPException(status_code=404, detail="Job not found")
+    return job_status
 
-@app.get("/api/health")
-async def health_check():
+@app.get("/api/jobs")
+async def list_jobs():
+    """Endpoint to list all jobs"""
     return {
-        "status": "healthy",
-        "service": "Storybook Generator API",
-        "timestamp": datetime.now().isoformat(),
-        "memory_usage_mb": monitor_memory(),
-        "models_loaded": list(model_cache.keys()),
-        "current_model": current_model_name,
-        "active_requests": len(active_requests)
+        "total_jobs": len(job_storage),
+        "jobs": {
+            job_id: {
+                "status": data["status"].value,
+                "progress": data["progress"],
+                "story_title": data["request"]["story_title"],
+                "created_at": data["created_at"]
+            }
+            for job_id, data in job_storage.items()
+        }
     }
 
-# Simple Gradio interface
-with gr.Blocks(title="Storybook Generator", theme="soft") as demo:
-    gr.Markdown("# 📚 Storybook Generator")
-    
-    with gr.Row():
-        story_title = gr.Textbox(label="Story Title", value="Test Story")
-        prompt_input = gr.Textbox(label="Scene Description", lines=3, value="A beautiful sunset over mountains with vibrant colors")
-        generate_btn = gr.Button("Generate Test Page")
-        output_image = gr.Image()
-        status = gr.Textbox()
-    
-    def generate_test_page(prompt, title):
-        try:
-            # Test with a simple generation
-            enhanced_prompt, negative_prompt = enhance_prompt(prompt, [], "childrens_book", 1)
-            
-            generator = torch.Generator(device="cpu")
-            generator.manual_seed(123)
-            
-            global current_pipe
-            image = current_pipe(
-                prompt=enhanced_prompt,
-                negative_prompt=negative_prompt,
-                num_inference_steps=20,
-                guidance_scale=7.0,
-                width=512,
-                height=512,
-                generator=generator
-            ).images[0]
-            
-            return image, f"✅ Generated: {enhanced_prompt}"
+
+# GRADIO INTERFACE (Optional - for manual testing)
+def create_gradio_interface():
+    """Create Gradio interface for manual testing"""
+    with gr.Blocks(title="High-Quality Storybook Generator", theme="soft") as demo:
+        gr.Markdown("# 🎨 High-Quality Storybook Generator")
+        gr.Markdown("Generate **studio-quality** storybook images with professional results")
+        
+        with gr.Row():
+            with gr.Column(scale=1):
+                gr.Markdown("### 🎯 Quality Settings")
+                
+                model_choice = gr.Dropdown(
+                    label="AI Model",
+                    choices=list(MODEL_CHOICES.keys()),
+                    value="dreamshaper-8",
+                    info="Choose the best model for your style"
+                )
+                
+                style_choice = gr.Dropdown(
+                    label="Art Style",
+                    choices=["childrens_book", "realistic", "fantasy", "anime"],
+                    value="childrens_book",
+                    info="Select the artistic style"
+                )
+                
+                story_title = gr.Textbox(
+                    label="Story Title",
+                    placeholder="Enter your story title..."
+                )
+                
+                scene_input = gr.Textbox(
+                    label="Scene Description",
+                    placeholder="Describe your scene in detail...",
+                    lines=3
+                )
+                
+                generate_btn = gr.Button("✨ Generate Premium Image", variant="primary")
             
-        except Exception as e:
-            return None, f"Error: {str(e)}"
-    
-    generate_btn.click(
-        fn=generate_test_page,
-        inputs=[prompt_input, story_title],
-        outputs=[output_image, status]
-    )
+            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=3)
+        
+        def generate_single_image(story_title, scene, model, style):
+            """Generate a single high-quality image for testing"""
+            try:
+                prompt, negative_prompt = enhance_prompt(scene, style)
+                image = generate_high_quality_image(prompt, model, style, negative_prompt)
+                return image, f"✅ High-quality image generated for: {story_title}"
+            except Exception as e:
+                return None, f"❌ Error: {str(e)}"
+        
+        generate_btn.click(
+            fn=generate_single_image,
+            inputs=[story_title, scene_input, model_choice, style_choice],
+            outputs=[image_output, status_output]
+        )
+        
+        return demo
 
-app = gr.mount_gradio_app(app, demo, path="/")
+# For Hugging Face Spaces deployment
+def get_app():
+    return app
 
 if __name__ == "__main__":
-    print("🚀 Starting Storybook Generator API...")
     import uvicorn
+    print("🚀 Starting High-Quality Storybook Generator API...")
+    print("🎨 Available models:")
+    for model in MODEL_CHOICES:
+        print(f"   - {model}")
+    
+    # Start both FastAPI and Gradio (if needed)
     uvicorn.run(app, host="0.0.0.0", port=7860)