Update app.py

app.py

@@ -15,12 +15,14 @@ os.environ["HF_HOME"] = "/tmp/huggingface"
 os.environ["TRANSFORMERS_CACHE"] = "/tmp/huggingface"
 os.environ["HF_HUB_CACHE"] = "/tmp/huggingface"
 
-# ✅ Load Public Image Generation Model (No Token Needed)
-IMAGE_MODEL = "stabilityai/sdxl-turbo"  # Fastest model for public access
+# ✅ Enable GPU if available
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# ✅ Load Public Image Generation Model
+IMAGE_MODEL = "runwayml/stable-diffusion-v1-5"  # ✅ Optimized for GPU
 pipeline = DiffusionPipeline.from_pretrained(
-    IMAGE_MODEL,
-    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
-).to("cuda" if torch.cuda.is_available() else "cpu")
+    IMAGE_MODEL, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+).to(device)
 
 # ✅ Define the input request format
 class StoryRequest(BaseModel):
@@ -41,30 +43,21 @@ def generate_story_questions_images(request: StoryRequest):
         story_text = story_result["story"]
         questions = story_result["questions"]
 
-        # ✅ Split the story into sentences for image generation
-        story_sentences = story_text.strip().split(". ")
-
-        # ✅ Generate an image for each sentence
-        images = []
-        for sentence in story_sentences:
-            if len(sentence) > 5:  # Avoid empty sentences
-                print(f"🖼️ Generating image for: {sentence}")
-                image = pipeline(prompt=sentence, num_inference_steps=5).images[0]
-
-                # Convert Image to Base64
-                img_byte_arr = io.BytesIO()
-                image.save(img_byte_arr, format="PNG")
-                img_base64 = base64.b64encode(img_byte_arr.getvalue()).decode("utf-8")
+        # ✅ Generate an image for the story theme
+        print(f"🖼️ Generating image for: {request.theme}")
+        image = pipeline(prompt=request.theme, num_inference_steps=5).images[0]
 
-                images.append({"sentence": sentence, "image": img_base64})
+        # Convert Image to Base64
+        img_byte_arr = io.BytesIO()
+        image.save(img_byte_arr, format="PNG")
+        img_base64 = base64.b64encode(img_byte_arr.getvalue()).decode("utf-8")
 
-        # ✅ Return the full response
         return {
             "theme": request.theme,
             "reading_level": request.reading_level,
             "story": story_text,
             "questions": questions,
-            "images": images,
+            "image": img_base64
         }
 
     except Exception as e: