Update app.py

app.py

@@ -1,8 +1,6 @@
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
-import torch
-from diffusers import StableDiffusionPipeline
-import story_generator  # ✅ Correctly importing the Story Generator
+import story_generator  # ✅ Import Story Generator
 
 app = FastAPI()
 
@@ -11,60 +9,29 @@ class StoryRequest(BaseModel):
     theme: str
     reading_level: str
 
-# ✅ Load AI Image Model (Cartoon-Style)
-CACHE_DIR = "/tmp/huggingface"
-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."""
+@app.post("/generate_story_and_questions")
+def generate_story_and_questions(request: StoryRequest):
+    """Generates a story and corresponding questions based on the given theme and reading level."""
     try:
         print(f"🎭 Generating story for theme: {request.theme} and level: {request.reading_level}")
 
-        # ✅ Generate the story
+        # ✅ Generate the story and questions
         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,
+            "story": story_result["story"],
             "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}")
+        print(f"❌ Error generating story and questions: {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 & Question Generation API! Use /generate_story_and_questions"}