Update app.py

app.py

@@ -1,67 +1,53 @@
 from fastapi import FastAPI
 from pydantic import BaseModel
-import uvicorn
 import torch
 from diffusers import StableDiffusionPipeline
 import story_generator  # Import Story Generator
 
 app = FastAPI()
 
-# ✅ Use a writable cache directory for image generation models
-CACHE_DIR = "/tmp/huggingface"
-torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
-
-# ✅ Load Image Model (Arcane-Diffusion for Cartoon-Style Images)
-print("Loading Cartoon Model...")
-image_model_id = "nitrosocke/Arcane-Diffusion"
-
-pipe = StableDiffusionPipeline.from_pretrained(
-    image_model_id,
-    cache_dir=CACHE_DIR,  # ✅ Store model in `/tmp`
-    torch_dtype=torch_dtype
-)
-pipe.to("cuda" if torch.cuda.is_available() else "cpu")
-print("Image Model Loaded!")
-
-# ✅ Define API Request Formats
+# ✅ Use a single request format
 class StoryRequest(BaseModel):
     theme: str
     reading_level: str
 
-class ImageRequest(BaseModel):
-    story: str
-
-# ✅ Story Generation Endpoint
-@app.post("/generate_story")
-def generate_story_endpoint(request: StoryRequest):
-    """Generate a story based on theme and reading level."""
-    result = story_generator.generate_story_and_questions(request.theme, request.reading_level)
-    return {
-        "theme": request.theme,
-        "reading_level": request.reading_level,
-        "story": result["story"],
-        "questions": result["questions"]
-    }
-
-# ✅ Image Generation Endpoint
-@app.post("/generate_image")
-def generate_image(request: ImageRequest):
-    """Generate a cartoon-style AI image based on the story text."""
-    cartoon_prompt = f"A colorful, cartoon-style illustration of: {request.story}, vibrant colors, highly detailed, storybook fantasy."
+# ✅ Load 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")
 
+@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])
-
-    # ✅ Generate Image
     image = pipe(cartoon_prompt, width=1024, height=1024).images[0]
 
-    # ✅ Save Image
+    # ✅ Save image
     image_path = "generated_image.png"
     image.save(image_path)
-
     print("Image Saved!")  
-    return {"message": "Image generated successfully!", "image_path": image_path}
 
-    @app.get("/")
+    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
+@app.get("/")
 def home():
     return {"message": "Welcome to the Story & Image Generation API! Use /generate_story_and_image"}
-