Update app.py

app.py

@@ -1,53 +1,70 @@
-from fastapi import FastAPI
+from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 import torch
 from diffusers import StableDiffusionPipeline
-import story_generator  # Import Story Generator
+import story_generator  # ✅ Correctly importing the Story Generator
 
 app = FastAPI()
 
-# ✅ Use a single request format
+# ✅ Define the input request format
 class StoryRequest(BaseModel):
     theme: str
     reading_level: str
 
-# ✅ Load Image Model (Cartoon-Style)
+# ✅ Load AI Image Model (Cartoon-Style)
 CACHE_DIR = "/tmp/huggingface"
-pipe = StableDiffusionPipeline.from_pretrained(
-    "nitrosocke/Arcane-Diffusion",
-    cache_dir=CACHE_DIR, 
-    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
-)
-pipe.to("cuda" if torch.cuda.is_available() else "cpu")
+try:
+    pipe = StableDiffusionPipeline.from_pretrained(
+        "nitrosocke/Arcane-Diffusion",
+        cache_dir=CACHE_DIR, 
+        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+    )
+    pipe.to("cuda" if torch.cuda.is_available() else "cpu")
+    print("✅ Image generation model loaded successfully.")
+except Exception as e:
+    print(f"❌ Failed to load image model: {e}")
+    pipe = None  # Fallback if model fails
 
 @app.post("/generate_story_and_image")
 def generate_story_and_image(request: StoryRequest):
     """Generates a story and an image based on the given theme and reading level."""
-    print(f"Generating story for theme: {request.theme} and level: {request.reading_level}")
-
-    # ✅ Generate story
-    story_result = story_generator.generate_story_and_questions(request.theme, request.reading_level)
-    story_text = story_result["story"]
-    
-    # ✅ Generate image
-    cartoon_prompt = f"A colorful, cartoon-style illustration of: {story_text}, vibrant colors, highly detailed, storybook fantasy."
-    print("Generating Image for:", cartoon_prompt[:100])
-    image = pipe(cartoon_prompt, width=1024, height=1024).images[0]
-
-    # ✅ Save image
-    image_path = "generated_image.png"
-    image.save(image_path)
-    print("Image Saved!")  
-
-    return {
-        "theme": request.theme,
-        "reading_level": request.reading_level,
-        "story": story_text,
-        "questions": story_result["questions"],
-        "image_path": image_path
-    }
-
-# ✅ Fix the "Not Found" issue for the root URL
+    try:
+        print(f"🎭 Generating story for theme: {request.theme} and level: {request.reading_level}")
+
+        # ✅ Generate the story
+        if not hasattr(story_generator, "generate_story_and_questions"):
+            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"]
+
+        # ✅ Generate the image
+        if not pipe:
+            raise HTTPException(status_code=500, detail="Image generation model failed to load.")
+
+        cartoon_prompt = f"A colorful, cartoon-style illustration of: {story_text}, vibrant colors, highly detailed, storybook fantasy."
+        print("🖌️ Generating Image for:", cartoon_prompt[:100])
+
+        image = pipe(cartoon_prompt, width=768, height=768).images[0]  # ✅ Slightly smaller for faster generation
+
+        # ✅ Save image
+        image_path = "/tmp/generated_image.png"
+        image.save(image_path)
+        print("✅ Image Saved!")
+
+        return {
+            "theme": request.theme,
+            "reading_level": request.reading_level,
+            "story": story_text,
+            "questions": story_result["questions"],
+            "image_url": f"https://your-api-url.com/generated_image.png"  # Replace with actual host
+        }
+
+    except Exception as e:
+        print(f"❌ Error generating story and image: {e}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+# ✅ Welcome message at root
 @app.get("/")
 def home():
-    return {"message": "Welcome to the Story & Image Generation API! Use /generate_story_and_image"}
+    return {"message": "🎉 Welcome to the Story & Image Generation API! Use /generate_story_and_image"}