Update app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 import torch
-from diffusers import StableDiffusionPipeline, EulerAncestralDiscreteScheduler
+from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler, DPMSolverMultistepScheduler
 from PIL import Image
 import io
 import requests
@@ -18,6 +18,8 @@ import random
 import gc
 import psutil
 import threading
+from transformers import CLIPTokenizer, CLIPTextModel
+import numpy as np
 
 # External OCI API URL
 OCI_API_BASE_URL = "https://yukee1992-oci-story-book.hf.space"
@@ -48,163 +50,268 @@ class StorybookRequest(BaseModel):
     story_title: str
     scenes: List[StoryScene]
     characters: List[CharacterDescription] = []
-    model_choice: str = "dreamshaper-8"
+    model_choice: str = "sdxl"
     style: str = "childrens_book"
 
-# HIGH-QUALITY MODEL SELECTION
+# MODEL SELECTION
 MODEL_CHOICES = {
+    "sdxl": "stabilityai/stable-diffusion-xl-base-1.0",
+    "sdxl-turbo": "stabilityai/sdxl-turbo",
     "dreamshaper-8": "lykon/dreamshaper-8",
     "realistic-vision": "SG161222/Realistic_Vision_V5.1", 
-    "anything-v5": "andite/anything-v5.0",
-    "openjourney": "prompthero/openjourney",
-    "sd-2.1": "stabilityai/stable-diffusion-2-1",
 }
 
-# GLOBAL MODEL CACHE - Load once, reuse forever
+# GLOBAL MODEL CACHE
 model_cache = {}
 current_model_name = None
 current_pipe = None
 
 # Character consistency tracking
 character_descriptions = {}
-character_seeds = {}  # Store seeds for consistent character generation
+character_seeds = {}
+
+# CLIP tokenizer for long prompt handling
+clip_tokenizer = None
+clip_model = None
+
+def initialize_clip():
+    """Initialize CLIP for long prompt processing"""
+    global clip_tokenizer, clip_model
+    try:
+        clip_tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14")
+        clip_model = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14")
+        print("✅ CLIP model loaded for long prompt processing")
+    except Exception as e:
+        print(f"❌ CLIP loading failed: {e}")
 
 # Memory monitoring function
 def monitor_memory():
-    """Monitor current memory usage"""
     try:
         process = psutil.Process()
-        memory_usage = process.memory_info().rss / 1024 / 1024  # MB
+        memory_usage = process.memory_info().rss / 1024 / 1024
         print(f"📊 Memory usage: {memory_usage:.2f} MB")
         return memory_usage
     except:
-        print("⚠️  Could not monitor memory (psutil not available)")
         return 0
 
-# Memory cleanup function
 def cleanup_memory():
-    """Clean up memory and cache"""
     gc.collect()
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
     print("🧹 Memory cleaned up")
 
-def load_model(model_name="dreamshaper-8"):
-    """Load model into global cache - runs only once per model"""
+def load_model(model_name="sdxl"):
     global model_cache, current_model_name, current_pipe
     
-    # Return cached model if already loaded
     if model_name in model_cache:
-        print(f"✅ Using cached model: {model_name}")
         current_pipe = model_cache[model_name]
         current_model_name = model_name
         return current_pipe
     
-    print(f"🔄 Loading model for the first time: {model_name}")
+    print(f"🔄 Loading model: {model_name}")
     try:
-        model_id = MODEL_CHOICES.get(model_name, "lykon/dreamshaper-8")
-        
-        pipe = StableDiffusionPipeline.from_pretrained(
-            model_id, 
-            torch_dtype=torch.float32,
-            safety_checker=None,
-            requires_safety_checker=False
-        )
+        if model_name in ["sdxl", "sdxl-turbo"]:
+            model_id = MODEL_CHOICES[model_name]
+            pipe = StableDiffusionXLPipeline.from_pretrained(
+                model_id,
+                torch_dtype=torch.float32,
+                use_safetensors=True,
+                safety_checker=None,
+                requires_safety_checker=False
+            )
+            pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
+        else:
+            model_id = MODEL_CHOICES.get(model_name, "lykon/dreamshaper-8")
+            pipe = StableDiffusionPipeline.from_pretrained(
+                model_id, 
+                torch_dtype=torch.float32,
+                safety_checker=None,
+                requires_safety_checker=False
+            )
+            pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
         
-        # Use better scheduler for quality
-        pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
         pipe = pipe.to("cpu")
-        
-        # Cache the model for future use
         model_cache[model_name] = pipe
         current_pipe = pipe
         current_model_name = model_name
         
-        print(f"✅ Model loaded and cached: {model_name}")
-        monitor_memory()
+        print(f"✅ Model loaded: {model_name}")
         return pipe
         
     except Exception as e:
         print(f"❌ Model loading failed: {e}")
-        # Fallback to SD 1.5
         pipe = StableDiffusionPipeline.from_pretrained(
             "runwayml/stable-diffusion-v1-5", 
-            torch_dtype=torch.float32,
-            safety_checker=None,
-            requires_safety_checker=False
+            torch_dtype=torch.float32
         ).to("cpu")
         model_cache[model_name] = pipe
-        current_pipe = pipe
         return pipe
 
-# Load the default model once at startup
+# Initialize CLIP and default model
 print("🚀 Initializing Storybook Generator...")
-current_pipe = load_model("dreamshaper-8")
-print("✅ Default model loaded and ready!")
-monitor_memory()
+initialize_clip()
+current_pipe = load_model("sdxl")
+print("✅ Models loaded and ready!")
 
-# PROFESSIONAL PROMPT ENGINEERING
-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"
-        ]
+# ADVANCED LONG PROMPT HANDLING
+def segment_long_prompt(long_prompt, max_tokens=75):
+    """
+    Split long prompt into meaningful segments using CLIP tokenization
+    and semantic analysis
+    """
+    if clip_tokenizer is None:
+        # Fallback: simple sentence splitting
+        sentences = [s.strip() for s in long_prompt.split('.') if s.strip()]
+        return sentences
+    
+    # Tokenize with CLIP to understand semantic boundaries
+    tokens = clip_tokenizer(long_prompt, return_tensors="pt", truncation=False)
+    token_count = tokens.input_ids.shape[1]
+    
+    if token_count <= max_tokens:
+        return [long_prompt]
+    
+    print(f"📝 Segmenting very long prompt: {token_count} tokens")
+    
+    # Split into sentences first
+    sentences = [s.strip() for s in long_prompt.split('.') if s.strip()]
+    segments = []
+    current_segment = ""
+    
+    for sentence in sentences:
+        test_segment = current_segment + ". " + sentence if current_segment else sentence
+        test_tokens = clip_tokenizer(test_segment, return_tensors="pt", truncation=False)
+        
+        if test_tokens.input_ids.shape[1] <= max_tokens:
+            current_segment = test_segment
+        else:
+            if current_segment:
+                segments.append(current_segment)
+            current_segment = sentence
+    
+    if current_segment:
+        segments.append(current_segment)
+    
+    return segments
+
+def create_prompt_hierarchy(full_prompt):
+    """
+    Create a hierarchical prompt structure with main focus and supporting details
+    """
+    segments = segment_long_prompt(full_prompt)
+    
+    if len(segments) == 1:
+        return full_prompt
+    
+    # The first segment is most important (main subject/action)
+    main_prompt = segments[0]
+    
+    # Remaining segments become supporting context with weights
+    supporting_context = ""
+    for i, segment in enumerate(segments[1:], 1):
+        weight = 1.3 - (i * 0.1)  # Decreasing weight for later segments
+        weight = max(0.8, min(1.5, weight))
+        supporting_context += f" ({segment}:{weight:.1f})"
+    
+    final_prompt = f"{main_prompt}.{supporting_context}. masterpiece, best quality, 4K"
+    return final_prompt
+
+def extract_key_phrases(prompt, max_phrases=10):
+    """
+    Extract the most important phrases from very long prompts
+    """
+    # Simple heuristic: nouns, adjectives, and verbs are important
+    words = prompt.split()
+    important_words = []
+    
+    # Prioritize words after colons, in parentheses, or quoted
+    for i, word in enumerate(words):
+        if (':' in word or '(' in word or '[' in word or 
+            word.isupper() or (i > 0 and words[i-1][-1] == ':')):
+            important_words.append(word)
+    
+    # Also take first few words of each sentence
+    sentences = prompt.split('.')
+    for sentence in sentences:
+        first_words = sentence.strip().split()[:3]
+        important_words.extend(first_words)
+    
+    # Remove duplicates and limit
+    important_words = list(set(important_words))[:max_phrases]
+    return " ".join(important_words)
+
+def enhance_prompt(scene_visual, characters, style="childrens_book", page_number=1):
+    """Create comprehensive prompt with NO length limits"""
+    
+    # Character context - include ALL details
+    character_context = ""
+    if characters:
+        char_descriptions = []
+        for char in characters:
+            if hasattr(char, 'description'):
+                char_descriptions.append(char.description)
+            elif isinstance(char, dict):
+                char_descriptions.append(char.get('description', ''))
+        character_context = " ".join(char_descriptions)
+        character_context = f"Character details: {character_context}."
+    
+    # Scene continuity context
+    continuity_context = f"Scene {page_number}, " if page_number > 1 else ""
+    
+    # Style templates
+    style_presets = {
+        "childrens_book": "children's book illustration, watercolor style, whimsical, charming, vibrant colors, soft lighting, storybook art, detailed backgrounds, cute characters, magical atmosphere",
+        "realistic": "photorealistic, professional photography, natural lighting, detailed, sharp focus, high resolution, realistic textures, studio quality, cinematic lighting",
+        "fantasy": "fantasy art, digital painting, magical, epic, concept art, dramatic lighting, mystical, otherworldly, detailed environments, heroic",
+        "anime": "anime style, Japanese animation, clean lines, vibrant colors, cel shading, detailed eyes, dynamic poses, manga style, professional animation"
     }
     
-    templates = style_templates.get(style, style_templates["childrens_book"])
-    style_prompt = templates[0]
+    style_prompt = style_presets.get(style, style_presets["childrens_book"])
     
-    enhanced = f"{style_prompt}, {prompt}"
+    # Build COMPREHENSIVE prompt with ALL details
+    full_prompt = f"""
+    {continuity_context}
+    {scene_visual}.
+    {character_context}
+    Art style: {style_prompt}.
+    Technical quality: masterpiece, best quality, 4K resolution, ultra detailed, 
+    professional artwork, award winning, trending on artstation, perfect composition,
+    ideal lighting, beautiful colors, no errors, perfect anatomy, consistent style
+    """
     
-    quality_boosters = [
-        "intricate details", "beautiful composition", "perfect lighting",
-        "professional artwork", "award winning", "trending on artstation"
-    ]
+    # Clean up the prompt
+    full_prompt = ' '.join(full_prompt.split())  # Remove extra whitespace
     
-    boosters = random.sample(quality_boosters, 2)
-    enhanced += ", " + ", ".join(boosters)
+    print(f"📝 Raw prompt length: {len(full_prompt.split())} words")
     
+    # Use hierarchical prompt creation for very long prompts
+    if len(full_prompt.split()) > 100:
+        optimized_prompt = create_prompt_hierarchy(full_prompt)
+    else:
+        optimized_prompt = full_prompt
+    
+    print(f"📝 Final prompt length: {len(optimized_prompt.split())} words")
+    
+    # Negative prompt
     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"
+        "blurry, low quality, ugly, deformed, poorly drawn, bad anatomy, "
+        "wrong anatomy, extra limb, missing limb, floating limbs, "
+        "disconnected limbs, mutation, mutated, disgusting, bad art, "
+        "beginner, amateur, distorted, watermark, signature, text, username, "
+        "multiple people, crowd, group, different characters, inconsistent features, "
+        "changed appearance, different face, altered features, low resolution, "
+        "jpeg artifacts, compression artifacts, noise, grain, out of focus"
     )
     
-    return enhanced, negative_prompt
+    return optimized_prompt, negative_prompt
 
 def save_complete_storybook_page(image, story_title, sequence_number, scene_text):
-    """Save image AND text to OCI with organized structure"""
     try:
-        # Convert image to bytes
         img_bytes = io.BytesIO()
         image.save(img_bytes, format='PNG')
         img_data = img_bytes.getvalue()
         
-        # Clean title for filenames
         clean_title = re.sub(r'[^a-zA-Z0-9_\-]', '', story_title.strip().replace(' ', '_'))
-        
-        # Create filenames
         image_filename = f"page_{sequence_number:03d}_{clean_title}.png"
         text_filename = f"page_{sequence_number:03d}_{clean_title}.txt"
         
@@ -229,112 +336,90 @@ def save_complete_storybook_page(image, story_title, sequence_number, scene_text
     except Exception as e:
         return f"❌ Save failed: {str(e)}"
 
-def enhance_with_character_context(scene_visual, story_title, characters):
-    """Add character descriptions to maintain consistency"""
-    if characters:
-        character_context = " ".join([f"{char.name}: {char.description}" for char in characters])
-        return f"Character descriptions: {character_context}. {scene_visual}"
-    return scene_visual
-
-def get_character_seed(story_title, character_name):
-    """Get consistent seed for character generation"""
+def get_character_seed(story_title, character_name, page_number):
     if story_title not in character_seeds:
         character_seeds[story_title] = {}
     
-    if character_name not in character_seeds[story_title]:
-        # Generate a stable seed based on character name and story title
-        seed_value = hash(f"{story_title}_{character_name}") % 1000000
-        character_seeds[story_title][character_name] = seed_value
-        print(f"🌱 Seed for {character_name}: {seed_value}")
+    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
     
-    return character_seeds[story_title][character_name]
+    return character_seeds[story_title][seed_key]
 
-def generate_storybook_page(scene_visual, story_title, sequence_number, scene_text, characters, model_choice="dreamshaper-8", style="childrens_book"):
-    """Generate a storybook page with character consistency"""
+def generate_storybook_page(scene_visual, story_title, sequence_number, scene_text, characters, model_choice="sdxl", style="childrens_book"):
     global current_pipe, current_model_name
     
     try:
-        # Switch model if different from current - BUT DON'T RELOAD UNLESS NECESSARY
         if model_choice != current_model_name:
-            print(f"🔄 Switching to model: {model_choice}")
-            current_pipe = load_model(model_choice)  # This uses cached version if available
-        
-        # ENHANCE PROMPT WITH CHARACTER CONTEXT
-        enhanced_visual = enhance_with_character_context(scene_visual, story_title, characters)
+            current_pipe = load_model(model_choice)
         
-        # Add scene continuity context
-        if sequence_number > 1:
-            enhanced_visual = f"Scene {sequence_number}, maintain character consistency from previous scenes. {enhanced_visual}"
+        enhanced_prompt, negative_prompt = enhance_prompt(
+            scene_visual, characters, style, sequence_number
+        )
         
-        enhanced_prompt, negative_prompt = enhance_prompt(enhanced_visual, style)
+        print(f"📖 Generating page {sequence_number}")
+        print(f"📝 Prompt preview: {enhanced_prompt[:150]}...")
         
-        print(f"📖 Generating page {sequence_number} for: {story_title}")
         if characters:
-            print(f"👤 Characters: {[char.name for char in characters]}")
+            char_names = []
+            for char in characters:
+                if hasattr(char, 'name'):
+                    char_names.append(char.name)
+                elif isinstance(char, dict):
+                    char_names.append(char.get('name', 'unknown'))
+            print(f"👤 Characters: {char_names}")
         
-        # Use consistent seed for character generation
         generator = torch.Generator(device="cpu")
+        
         if characters:
-            # Use seed from main character for consistency
-            main_char_seed = get_character_seed(story_title, characters[0].name)
+            first_char = characters[0]
+            char_name = first_char.name if hasattr(first_char, 'name') else first_char.get('name', 'unknown')
+            main_char_seed = get_character_seed(story_title, char_name, sequence_number)
             generator.manual_seed(main_char_seed)
-            print(f"🌱 Using seed {main_char_seed} for character consistency")
         else:
-            seed = int(time.time())
-            generator.manual_seed(seed)
-            print(f"🌱 Using timestamp seed {seed}")
+            scene_seed = hash(f"{story_title}_{sequence_number}") % 1000000
+            generator.manual_seed(scene_seed)
         
-        # Generate high-quality image - USE THE GLOBAL current_pipe
+        # Generate with SDXL which handles long prompts better
         image = current_pipe(
             prompt=enhanced_prompt,
             negative_prompt=negative_prompt,
-            num_inference_steps=30,
-            guidance_scale=8.5,
+            num_inference_steps=40,  # More steps for better detail
+            guidance_scale=7.0,
             width=768,
             height=768,
             generator=generator
         ).images[0]
         
-        # Save both image and text
         save_status = save_complete_storybook_page(image, story_title, sequence_number, scene_text)
-        
         return image, save_status
         
     except Exception as e:
         return None, f"❌ Generation failed: {str(e)}"
 
-def batch_generate_complete_storybook(story_title, scenes_data, characters, model_choice="dreamshaper-8", style="childrens_book"):
-    """Generate complete storybook with memory management"""
+def batch_generate_complete_storybook(story_title, scenes_data, characters, model_choice="sdxl", style="childrens_book"):
     global character_descriptions, current_pipe
     
     results = []
     status_messages = []
     
-    print(f"📚 Starting batch generation for: {story_title}")
-    print(f"📖 Total pages: {len(scenes_data)}")
+    print(f"📚 Starting batch generation: {story_title}")
+    print(f"📖 Pages: {len(scenes_data)}")
     print(f"👤 Characters: {len(characters)}")
-    print(f"🎨 Using model: {model_choice}")
     
-    # Initial memory check
-    initial_memory = monitor_memory()
-    
-    # Store character descriptions for this story
     if characters:
         character_descriptions[story_title] = characters
-        print(f"✅ Character context stored for {story_title}")
     
-    # Load model once at the beginning
-    print(f"🔧 Loading model for this storybook...")
     current_pipe = load_model(model_choice)
-    
     start_time = time.time()
     
     for i, scene_data in enumerate(scenes_data, 1):
         try:
-            # Clean memory every 2 pages
             if i % 2 == 0:
                 cleanup_memory()
-                monitor_memory()
             
             scene_visual = scene_data.get('visual', '')
             scene_text = scene_data.get('text', '')
@@ -348,43 +433,34 @@ def batch_generate_complete_storybook(story_title, scenes_data, characters, mode
                 results.append((f"Page {i}", image, scene_text))
             status_messages.append(f"Page {i}: {status}")
             
+            if i < len(scenes_data):
+                time.sleep(2)
+            
         except Exception as e:
             error_msg = f"❌ Failed page {i}: {str(e)}"
             print(error_msg)
             status_messages.append(error_msg)
-            # Continue with next page instead of stopping
     
     total_time = time.time() - start_time
-    final_memory = monitor_memory()
-    
-    print(f"✅ Batch generation completed in {total_time:.2f} seconds")
-    print(f"📊 Memory delta: {final_memory - initial_memory:.2f} MB")
+    print(f"✅ Batch completed in {total_time:.2f} seconds")
     
     return results, "\n".join(status_messages)
 
-# FastAPI endpoint for n8n
+# FastAPI endpoint
 @app.post("/api/generate-storybook")
 async def api_generate_storybook(request: StorybookRequest):
-    """API endpoint for n8n automation - OPTIMIZED with character consistency"""
     try:
-        print(f"📚 Received storybook request: {request.story_title}")
-        print(f"📖 Pages to generate: {len(request.scenes)}")
-        print(f"👤 Characters received: {len(request.characters)}")
-        
-        if request.characters:
-            for char in request.characters:
-                print(f"   - {char.name}: {char.description[:50]}...")
+        print(f"📚 Received request: {request.story_title}")
+        print(f"📖 Pages: {len(request.scenes)}")
         
         start_time = time.time()
-        
-        # Convert to scene data format
         scenes_data = [{"visual": scene.visual, "text": scene.text} for scene in request.scenes]
+        characters_dict = [char.dict() for char in request.characters]
         
-        # Generate storybook (model loads only once)
         results, status = batch_generate_complete_storybook(
             request.story_title,
             scenes_data,
-            request.characters,
+            characters_dict,
             request.model_choice,
             request.style
         )
@@ -412,113 +488,24 @@ async def api_generate_storybook(request: StorybookRequest):
     except Exception as e:
         error_msg = f"Storybook generation failed: {str(e)}"
         print(f"❌ {error_msg}")
-        import traceback
-        traceback.print_exc()
         raise HTTPException(status_code=500, detail=error_msg)
 
-# Async processing endpoint for large batches
-@app.post("/api/generate-storybook-async")
-async def api_generate_storybook_async(request: StorybookRequest, background_tasks: BackgroundTasks):
-    """Async endpoint that processes images in background with memory management"""
-    try:
-        # Store the request and return immediate response
-        request_id = f"{request.story_title}_{int(time.time())}"
-        
-        # Start background task
-        background_tasks.add_task(
-            process_storybook_async, 
-            request_id, 
-            request.dict()
-        )
-        
-        return {
-            "status": "processing",
-            "request_id": request_id,
-            "message": f"Started processing {len(request.scenes)} pages for '{request.story_title}'",
-            "estimated_time": f"Approximately {len(request.scenes) * 45} seconds"
-        }
-        
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
-
-def process_storybook_async(request_id, request_data):
-    """Background task for async processing with memory management"""
-    try:
-        print(f"🔧 Starting async processing for request: {request_id}")
-        print(f"📖 Pages to process: {len(request_data['scenes'])}")
-        
-        # Convert dictionary back to proper objects for character handling
-        characters = []
-        if 'characters' in request_data and request_data['characters']:
-            # Convert character dicts back to CharacterDescription objects
-            for char_dict in request_data['characters']:
-                characters.append(CharacterDescription(**char_dict))
-        
-        # Initial memory check
-        initial_memory = monitor_memory()
-        
-        for i, scene in enumerate(request_data['scenes']):
-            try:
-                print(f"🔄 Processing page {i+1}/{len(request_data['scenes'])} for request {request_id}")
-                
-                # Generate single page - pass the converted character objects
-                image, status = generate_storybook_page(
-                    scene['visual'],
-                    request_data['story_title'],
-                    i+1,
-                    scene['text'],
-                    characters,  # Pass the converted objects, not raw dicts
-                    request_data.get('model_choice', 'dreamshaper-8'),
-                    request_data.get('style', 'childrens_book')
-                )
-                
-                print(f"✅ Page {i+1} completed: {status}")
-                
-                # Clean memory after each page
-                cleanup_memory()
-                current_memory = monitor_memory()
-                
-                # Add delay between pages to prevent overload
-                if i < len(request_data['scenes']) - 1:  # Don't sleep after last page
-                    sleep_time = 5  # 5 second delay between pages
-                    print(f"⏳ Waiting {sleep_time} seconds before next page...")
-                    time.sleep(sleep_time)
-                    
-            except Exception as e:
-                error_msg = f"❌ Failed page {i+1}: {str(e)}"
-                print(error_msg)
-                # Continue with next page
-                continue
-        
-        final_memory = monitor_memory()
-        print(f"✅ Completed async processing for {request_id}")
-        print(f"📊 Total memory change: {final_memory - initial_memory:.2f} MB")
-        
-    except Exception as e:
-        print(f"❌ Async processing failed for {request_id}: {e}")
-
-# Health check endpoint with memory info
 @app.get("/api/health")
 async def health_check():
-    memory_info = monitor_memory()
     return {
         "status": "healthy",
         "service": "Storybook Generator API",
         "timestamp": datetime.now().isoformat(),
-        "memory_usage_mb": round(memory_info, 2),
+        "memory_usage_mb": monitor_memory(),
         "models_loaded": list(model_cache.keys()),
-        "current_model": current_model_name,
-        "cached_models_count": len(model_cache),
-        "stories_tracked": len(character_descriptions)
+        "current_model": current_model_name
     }
 
-# Gradio Interface Functions
+# Gradio Interface
 def generate_single_page(prompt, story_title, scene_text, model_choice, style):
-    """Generate a single page for Gradio interface"""
     if not prompt or not story_title:
         return None, "❌ Please enter both scene description and story title"
     
-    # Ensure model is loaded
     global current_pipe
     if current_model_name != model_choice:
         current_pipe = load_model(model_choice)
@@ -528,25 +515,18 @@ def generate_single_page(prompt, story_title, scene_text, model_choice, style):
     )
     return image, status
 
-# Create the Gradio interface
 with gr.Blocks(title="Storybook Generator", theme="soft") as demo:
     gr.Markdown("# 📚 Storybook Generator")
     gr.Markdown("Create beautiful storybooks with consistent characters")
     
     with gr.Row():
         with gr.Column(scale=1):
-            story_title_input = gr.Textbox(
-                label="Story Title",
-                placeholder="Enter your story title...",
-                lines=1
-            )
-            
+            story_title_input = gr.Textbox(label="Story Title", lines=1)
             model_choice = gr.Dropdown(
                 label="AI Model",
                 choices=list(MODEL_CHOICES.keys()),
-                value="dreamshaper-8"
+                value="sdxl"
             )
-            
             style_choice = gr.Dropdown(
                 label="Art Style",
                 choices=["childrens_book", "realistic", "fantasy", "anime"],
@@ -554,18 +534,8 @@ with gr.Blocks(title="Storybook Generator", theme="soft") as demo:
             )
         
         with gr.Column(scale=2):
-            prompt_input = gr.Textbox(
-                label="Visual Description",
-                placeholder="Describe the scene for image generation...",
-                lines=3
-            )
-            
-            text_input = gr.Textbox(
-                label="Story Text (Optional)",
-                placeholder="Enter the story text for this page...",
-                lines=2
-            )
-            
+            prompt_input = gr.Textbox(label="Visual Description", lines=5)
+            text_input = gr.Textbox(label="Story Text (Optional)", lines=2)
             generate_btn = gr.Button("✨ Generate Single Page", variant="primary")
             image_output = gr.Image(label="Generated Page", height=400)
             status_output = gr.Textbox(label="Status", interactive=False)
@@ -576,21 +546,9 @@ with gr.Blocks(title="Storybook Generator", theme="soft") as demo:
         outputs=[image_output, status_output]
     )
 
-# Mount Gradio app to FastAPI
 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...")
-    print("📚 Available models:", list(MODEL_CHOICES.keys()))
-    print("🌐 API endpoints:")
-    print("   - POST /api/generate-storybook")
-    print("   - POST /api/generate-storybook-async (for large batches)")
-    print("   - GET  /api/health")
-    print("   - GET  / (Gradio UI)")
-    
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=7860)