Update app.py

app.py

@@ -1,4 +1,5 @@
 import os
+import time
 import torch
 from fastapi import FastAPI, HTTPException
 from fastapi.responses import JSONResponse
@@ -13,7 +14,7 @@ app = FastAPI()
 
 # Set Hugging Face cache directories
 os.environ["HF_HOME"] = "/tmp/huggingface"
-os.environ["TRANSFORMERS_CACHE"] = "/tmp/huggingface"
+os.environ["TRANSFORMERS_CACHE"] = "/tmp/huggingface"  # Deprecated but still used for now
 os.environ["HF_HUB_CACHE"] = "/tmp/huggingface"
 
 # Enable GPU if available
@@ -31,24 +32,64 @@ class StoryRequest(BaseModel):
     theme: str
     reading_level: str
 
+def generate_images(prompts, max_retries=3, delay=2):
+    """
+    Attempts to generate images in batch. If an error related to 
+    "index 16 is out of bounds" occurs, it retries for up to max_retries.
+    If all attempts fail, it falls back to generating images sequentially.
+    """
+    for attempt in range(max_retries):
+        try:
+            print(f"Batched image generation attempt {attempt+1}...")
+            results = pipeline(
+                prompt=prompts,
+                num_inference_steps=15,
+                height=768,
+                width=768
+            ).images
+            return results
+        except Exception as e:
+            if "index 16 is out of bounds" in str(e):
+                print(f"Attempt {attempt+1} failed with error: {e}")
+                time.sleep(delay)
+            else:
+                raise e
+    # Fallback to sequential generation
+    print("Falling back to sequential image generation...")
+    images = []
+    for i, prompt in enumerate(prompts):
+        try:
+            print(f"Sequential generation for prompt {i+1}...")
+            image = pipeline(
+                prompt=prompt,
+                num_inference_steps=15,
+                height=768,
+                width=768
+            ).images[0]
+            images.append(image)
+        except Exception as e:
+            print(f"Error in sequential generation for prompt {i+1}: {e}")
+            raise e
+    return images
+
 @app.post("/generate_story_questions_images")
 def generate_story_questions_images(request: StoryRequest):
-    """Generates a story, dynamic questions, and cartoonish storybook images."""
+    """
+    Generates a story, dynamic questions, and cartoonish storybook images.
+    """
     try:
         print(f"🎭 Generating story for theme: {request.theme} and level: {request.reading_level}")
-        
-        # Generate story and questions
+        # Generate story and questions using the story_generator module
         story_result = story_generator.generate_story_and_questions(request.theme, request.reading_level)
         story_text = story_result.get("story", "").strip()
         questions = story_result.get("questions", "").strip()
-        
         if not story_text:
             raise HTTPException(status_code=500, detail="Story generation failed.")
         
         # Split the story into up to 6 paragraphs
         paragraphs = [p.strip() for p in story_text.split("\n") if p.strip()][:6]
         
-        # Batch image generation:
+        # Build a list of prompts for batched image generation
         prompts = [
             (
                 f"Children's storybook illustration of: {p}. "
@@ -59,9 +100,10 @@ def generate_story_questions_images(request: StoryRequest):
         ]
         print(f"Generating images for {len(prompts)} paragraphs concurrently...")
         
-        # Single batched call to generate images concurrently
-        results = pipeline(prompt=prompts, num_inference_steps=15, height=768, width=768).images
+        # Use the retry mechanism for image generation
+        results = generate_images(prompts, max_retries=3, delay=2)
         
+        # Convert each generated image to Base64
         images = []
         for image in results:
             img_byte_arr = io.BytesIO()
@@ -77,11 +119,10 @@ def generate_story_questions_images(request: StoryRequest):
             "questions": questions,
             "images": images
         })
-    
     except Exception as e:
         print(f"❌ Error generating story/questions/images: {e}")
         raise HTTPException(status_code=500, detail=str(e))
 
 @app.get("/")
 def home():
-    return {"message": "🎉 Welcome to the Story, Question & Image API!"}
+    return {"message": "🎉 Welcome to the Story, Question & Image API!"}