Update app.py

app.py

@@ -3,7 +3,7 @@ import torch
 from fastapi import FastAPI, HTTPException
 from fastapi.responses import JSONResponse
 from pydantic import BaseModel
-import story_generator  # Make sure story_generator.py is correctly uploaded
+import story_generator
 from diffusers import DiffusionPipeline
 from PIL import Image
 import io
@@ -11,7 +11,7 @@ import base64
 
 app = FastAPI()
 
-# ✅ Hugging Face cache directory
+# ✅ Set Hugging Face cache directory
 os.environ["HF_HOME"] = "/tmp/huggingface"
 os.environ["TRANSFORMERS_CACHE"] = "/tmp/huggingface"
 os.environ["HF_HUB_CACHE"] = "/tmp/huggingface"
@@ -21,47 +21,47 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 
 # ✅ Load image generation model
 IMAGE_MODEL = "lykon/dreamshaper-8"
-pipeline = DiffusionPipeline.from_pretrained(IMAGE_MODEL, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32).to(device)
+pipeline = DiffusionPipeline.from_pretrained(
+    IMAGE_MODEL, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+).to(device)
+
 
 # ✅ Define request format
 class StoryRequest(BaseModel):
     theme: str
     reading_level: str
 
+
 @app.post("/generate_story_questions_images")
 def generate_story_questions_images(request: StoryRequest):
-    """Generates a story, questions, and cartoonish images."""
+    """Generates a story, questions, and cartoonish storybook images."""
     try:
         print(f"🎭 Generating story for theme: {request.theme} and level: {request.reading_level}")
 
-        # ✅ Generate the story and questions
+        # ✅ Generate the story
         if not hasattr(story_generator, "generate_story_and_questions"):
-            raise HTTPException(status_code=500, detail="Story generator function not found.")
+            raise HTTPException(status_code=500, detail="Story generation function not found in story_generator.py")
 
         story_result = story_generator.generate_story_and_questions(request.theme, request.reading_level)
-        story_text = story_result["story"]
-        questions = story_result["questions"]
-
-        # ✅ Process Story (Ensure it’s correctly structured)
-        if not isinstance(story_text, str) or len(story_text) < 50:
-            raise HTTPException(status_code=500, detail="Generated story is too short or invalid.")
+        story_text = story_result.get("story", "").strip()
+        questions = story_result.get("questions", "").strip()
 
-        # ✅ Process Questions (Ensure they are valid)
-        if not isinstance(questions, list) or len(questions) < 3:
+        # 🚨 Error Handling: Ensure valid story & questions are generated
+        if not story_text:
+            raise HTTPException(status_code=500, detail="Story generation failed.")
+        if not questions or "?" not in questions:
             raise HTTPException(status_code=500, detail="Question generation failed.")
 
-        # ✅ Generate cartoon-style images (for each paragraph)
-        paragraphs = [p.strip() for p in story_text.split("\n") if p.strip()]
+        # ✅ Split the story into paragraphs (max 6)
+        paragraphs = [p.strip() for p in story_text.split("\n") if p.strip()][:6]  # Only take up to 6 paragraphs
+
+        # ✅ Generate high-quality cartoonish images
         images = []
-        
         for i, paragraph in enumerate(paragraphs):
-            if i >= 6:  # ✅ Limit image generation to first 6 paragraphs (to save time)
-                break
-
-            prompt = f"Children's storybook illustration of: {paragraph}. Soft pastel colors, hand-drawn style, warm lighting, fantasy setting, watercolor texture."
-            print(f"🖼️ Generating better cartoon-style image for paragraph {i+1}...")
+            prompt = f"Children's storybook illustration of: {paragraph}. Soft pastel colors, hand-drawn style, friendly characters, warm lighting, fantasy setting, watercolor texture, storybook illustration, beautiful composition."
+            print(f"🖼 Generating cartoon-style image for paragraph {i+1}...")
 
-            # ✅ Reduce inference steps (from 40 → 15 for speed)
+            # ✅ Reduce inference steps to speed it up (from 40 → 15)
             image = pipeline(prompt=prompt, num_inference_steps=15, height=768, width=768).images[0]
 
             # ✅ Convert image to Base64
@@ -81,10 +81,11 @@ def generate_story_questions_images(request: StoryRequest):
         })
 
     except Exception as e:
-        print(f"❌ Error generating content: {e}")
+        print(f"❌ Error generating story/questions/images: {e}")
         raise HTTPException(status_code=500, detail=str(e))
 
+
 # ✅ Welcome message
 @app.get("/")
 def home():
-    return {"message": "🎉 Welcome to the Story, Question & Image API!"}
+    return {"message": "🎉 Welcome to the Story, Question & Image API!"}