Rename app (2).py to app.py

app (2).py → app.py

@@ -1,40 +1,51 @@
 from fastapi import FastAPI
 from pydantic import BaseModel
-import uvicorn
-from diffusers import StableDiffusionPipeline
 import torch
+from diffusers import StableDiffusionPipeline
+import story_generator  # Import Story Generator
 
 app = FastAPI()
 
-# ✅ Use `nitrosocke/Arcane-Diffusion` for high-quality cartoon-style images
-model_id = "nitrosocke/Arcane-Diffusion"
-
-# ✅ Check for GPU, else use CPU (Important for Render!)
+# ✅ Load Image Model (Arcane-Diffusion for Cartoon-Style Images)
+print("Loading Cartoon Model...")
+image_model_id = "nitrosocke/Arcane-Diffusion"
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# ✅ Load the AI Model ONCE (so it doesn’t reload every request)
-print("Loading Cartoon Model...")
 pipe = StableDiffusionPipeline.from_pretrained(
-    model_id, torch_dtype=torch.float16 if device == "cuda" else torch.float32
+    image_model_id, torch_dtype=torch.float16 if device == "cuda" else torch.float32
 )
 pipe.to(device)
-print("Model Loaded!")
+print("Image Model Loaded!")
+
+# ✅ Define API Request Formats
+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 high-quality cartoon-style image based on the story text."""
-    
-    cartoon_prompt = (
-        f"A colorful, cartoon-style illustration of: {request.story}, "
-        "storybook fantasy world, highly detailed, vibrant colors, soft edges, clear lighting"
-    )
+    """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."
 
     print("Generating Image for:", cartoon_prompt[:100])
 
-    # ✅ Generate Image (Directly in HD 1024x1024)
+    # ✅ Generate Image
     image = pipe(cartoon_prompt, width=1024, height=1024).images[0]
 
     # ✅ Save Image
@@ -42,8 +53,4 @@ def generate_image(request: ImageRequest):
     image.save(image_path)
 
     print("Image Saved!")  
-
     return {"message": "Image generated successfully!", "image_path": image_path}
-
-if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=8000)