Update app.py

app.py

@@ -1,190 +1,87 @@
+import os
 import torch
-from transformers import GPT2Tokenizer, GPT2LMHeadModel
-import random
-import re
+from fastapi import FastAPI, HTTPException
+from fastapi.responses import JSONResponse
+from pydantic import BaseModel
 import story_generator
+from diffusers import DiffusionPipeline
+from PIL import Image
+import io
+import base64
 
-# Set Hugging Face cache directory
-CACHE_DIR = "/tmp/huggingface"
+app = FastAPI()
 
-# ------------------------
-# Load Story Generation Model
-# ------------------------
-STORY_MODEL_NAME = "abdalraheemdmd/story-api"
-story_tokenizer = GPT2Tokenizer.from_pretrained(STORY_MODEL_NAME, cache_dir=CACHE_DIR)
-story_model = GPT2LMHeadModel.from_pretrained(STORY_MODEL_NAME, cache_dir=CACHE_DIR)
+# Set Hugging Face cache directories
+os.environ["HF_HOME"] = "/tmp/huggingface"
+os.environ["TRANSFORMERS_CACHE"] = "/tmp/huggingface"
+os.environ["HF_HUB_CACHE"] = "/tmp/huggingface"
 
-# ------------------------
-# Load Question Generation Model
-# ------------------------
-QUESTION_MODEL_NAME = "abdalraheemdmd/question-gene"
-question_tokenizer = GPT2Tokenizer.from_pretrained(QUESTION_MODEL_NAME, cache_dir=CACHE_DIR)
-question_model = GPT2LMHeadModel.from_pretrained(QUESTION_MODEL_NAME, cache_dir=CACHE_DIR)
+# Enable GPU if available
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Ensure tokenizers have a pad token
-if story_tokenizer.pad_token_id is None:
-    story_tokenizer.pad_token_id = story_tokenizer.eos_token_id
-if question_tokenizer.pad_token_id is None:
-    question_tokenizer.pad_token_id = question_tokenizer.eos_token_id
+# Load image generation model
+IMAGE_MODEL = "lykon/dreamshaper-8"
+pipeline = DiffusionPipeline.from_pretrained(
+    IMAGE_MODEL,
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+).to(device)
 
-def generate_story(theme, reading_level, max_new_tokens=400, temperature=0.7):
-    """Generates a story based on the provided theme and reading level."""
-    prompt = f"A {reading_level} story about {theme}:"
-    input_ids = story_tokenizer(prompt, return_tensors="pt").input_ids
-    with torch.no_grad():
-        output = story_model.generate(
-            input_ids,
-            max_new_tokens=max_new_tokens,
-            temperature=temperature,
-            top_k=20,
-            top_p=0.7,
-            do_sample=True,
-            early_stopping=True,
-            pad_token_id=story_tokenizer.pad_token_id,
-            eos_token_id=story_tokenizer.eos_token_id,
-            attention_mask=input_ids.ne(story_tokenizer.pad_token_id)
-        )
-    return story_tokenizer.decode(output[0], skip_special_tokens=True)
+# Define request schema
+class StoryRequest(BaseModel):
+    theme: str
+    reading_level: str
 
-def extract_protagonist(story):
-    """
-    Attempts to extract the protagonist from the first sentence by searching for the pattern "named <Name>".
-    Returns the first matched name, if available.
-    """
-    sentences = re.split(r'\.|\n', story)
-    if sentences:
-        m = re.search(r"named\s+([A-Z][a-z]+)", sentences[0])
-        if m:
-            return m.group(1)
-    return None
-
-def extract_characters(story):
-    """
-    Extracts potential character names from the story using a frequency count on capitalized words.
-    Filters out common stopwords so that the most frequently mentioned name is likely the main character.
-    """
-    words = re.findall(r'\b[A-Z][a-zA-Z]+\b', story)
-    stopwords = {"The", "A", "An", "And", "But", "Suddenly", "Quickly", "However", "Well",
-                 "They", "I", "He", "She", "It", "When", "Where", "Dr", "Mr"}
-    filtered = [w for w in words if w not in stopwords and len(w) > 2]
-    if not filtered:
-        return []
-    freq = {}
-    for word in filtered:
-        freq[word] = freq.get(word, 0) + 1
-    sorted_chars = sorted(freq.items(), key=lambda x: x[1], reverse=True)
-    return [item[0] for item in sorted_chars]
-
-def extract_themes(story):
-    """Extracts themes from the story based on keyword matching."""
-    themes = []
-    story_lower = story.lower()
-    if "space" in story_lower:
-        themes.append("space")
-    if "adventure" in story_lower:
-        themes.append("adventure")
-    if "friend" in story_lower:
-        themes.append("friendship")
-    if "learn" in story_lower or "lesson" in story_lower:
-        themes.append("learning")
-    return themes
-
-def extract_lesson(story):
-    """
-    Attempts to extract a lesson or moral from the story by finding sentences
-    containing keywords like "learn" or "lesson". Returns the last matching sentence.
-    """
-    sentences = re.split(r'\.|\n', story)
-    lesson_sentences = [
-        s.strip() for s in sentences
-        if ("learn" in s.lower() or "lesson" in s.lower()) and len(s.strip()) > 20
-    ]
-    if lesson_sentences:
-        return lesson_sentences[-1]
-    else:
-        return "No explicit lesson found."
-
-def format_question(question_prompt, correct_answer, distractors):
-    """
-    Combines the correct answer with three distractors, shuffles the options,
-    and formats the question as a multiple-choice question.
-    """
-    # Ensure exactly 3 distractors are available
-    if len(distractors) < 3:
-        default_distractors = ["Option X", "Option Y", "Option Z"]
-        while len(distractors) < 3:
-            distractors.append(default_distractors[len(distractors) % len(default_distractors)])
-    else:
-        distractors = random.sample(distractors, 3)
-    options = distractors + [correct_answer]
-    random.shuffle(options)
-    letters = ["A", "B", "C", "D"]
-    correct_letter = letters[options.index(correct_answer)]
-    options_text = "\n".join(f"{letters[i]}) {option}" for i, option in enumerate(options))
-    question_text = f"{question_prompt}\n{options_text}\nCorrect Answer: {correct_letter}"
-    return question_text
-
-def dynamic_fallback_questions(story):
-    """
-    Generates three multiple-choice questions based on dynamic story content.
-    Each question uses a randomly chosen template and shuffles its options.
-    """
-    protagonist = extract_protagonist(story)
-    characters = extract_characters(story)
-    themes = extract_themes(story)
-    lesson = extract_lesson(story)
+@app.post("/generate_story_questions_images")
+def generate_story_questions_images(request: StoryRequest):
+    """Generates a story, dynamic questions, and cartoonish storybook images."""
+    try:
+        print(f"🎭 Generating story for theme: {request.theme} and level: {request.reading_level}")
+        
+        # Generate story and questions
+        story_result = story_generator.generate_story_and_questions(request.theme, request.reading_level)
+        story_text = story_result.get("story", "").strip()
+        questions = story_result.get("questions", "").strip()
+        
+        if not story_text:
+            raise HTTPException(status_code=500, detail="Story generation failed.")
+        
+        # Split the story into up to 6 paragraphs
+        paragraphs = [p.strip() for p in story_text.split("\n") if p.strip()][:6]
+        
+        # Batch image generation:
+        prompts = [
+            (
+                f"Children's storybook illustration of: {p}. "
+                "Soft pastel colors, hand-drawn style, friendly characters, warm lighting, "
+                "fantasy setting, watercolor texture, storybook illustration, beautiful composition."
+            )
+            for p in paragraphs
+        ]
+        print(f"Generating images for {len(prompts)} paragraphs concurrently...")
+        
+        # Single batched call to generate images concurrently
+        results = pipeline(prompt=prompts, num_inference_steps=15, height=768, width=768).images
+        
+        images = []
+        for image in results:
+            img_byte_arr = io.BytesIO()
+            image.save(img_byte_arr, format="PNG")
+            img_byte_arr.seek(0)
+            base64_image = base64.b64encode(img_byte_arr.getvalue()).decode("utf-8")
+            images.append(base64_image)
+        
+        return JSONResponse(content={
+            "theme": request.theme,
+            "reading_level": request.reading_level,
+            "story": story_text,
+            "questions": questions,
+            "images": images
+        })
     
-    # --- Question 1: Theme ---
-    theme_templates = [
-        "What is the main theme of the story?",
-        "Which theme best represents the narrative?",
-        "What subject is central to the story?"
-    ]
-    q1_prompt = random.choice(theme_templates)
-    correct_theme = " and ".join(themes) if themes else "learning"
-    q1_distractors = ["sports and competition", "cooking and baking", "weather and seasons", "technology and innovation"]
-    q1 = format_question(q1_prompt, correct_theme, q1_distractors)
-    
-    # --- Question 2: Primary Character ---
-    character_templates = [
-        "Who is the primary character in the story?",
-        "Which character drives the main action in the narrative?",
-        "Who is the central figure in the story?"
-    ]
-    q2_prompt = random.choice(character_templates)
-    if protagonist:
-        correct_character = protagonist
-    elif characters:
-        correct_character = characters[0]
-    else:
-        correct_character = "The main character"
-    q2_distractors = ["a mysterious stranger", "an unknown visitor", "a supporting character", "a sidekick"]
-    q2 = format_question(q2_prompt, correct_character, q2_distractors)
-    
-    # --- Question 3: Lesson/Moral ---
-    lesson_templates = [
-        "What lesson did the characters learn by the end of the story?",
-        "What moral can be inferred from the narrative?",
-        "What is the key takeaway from the story?"
-    ]
-    q3_prompt = random.choice(lesson_templates)
-    if lesson and lesson != "No explicit lesson found.":
-        correct_lesson = lesson  # full sentence without truncation
-    else:
-        correct_lesson = "understanding and growth"
-    q3_distractors = ["always be silent", "never try new things", "do nothing", "ignore opportunities"]
-    q3 = format_question(q3_prompt, correct_lesson, q3_distractors)
-    
-    return f"{q1}\n\n{q2}\n\n{q3}"
-
-def generate_story_and_questions(theme, reading_level):
-    """
-    Generates a story using the story generation model and then creates dynamic,
-    multiple-choice questions based on that story.
-    """
-    story = generate_story(theme, reading_level)
-    questions = dynamic_fallback_questions(story)
-    return {"story": story, "questions": questions}
+    except Exception as e:
+        print(f"❌ Error generating story/questions/images: {e}")
+        raise HTTPException(status_code=500, detail=str(e))
 
-# Alias for backward compatibility
-create_fallback_questions = dynamic_fallback_questions
+@app.get("/")
+def home():
+    return {"message": "🎉 Welcome to the Story, Question & Image API!"}