Update backend.py

backend.py

@@ -1,72 +1,84 @@
 import os
+import time
+import torch
 from flask import Flask, request, jsonify
 from flask_cors import CORS
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import torch
-
-# Path to your local model directory (update as needed)
-MODEL_PATH = "ibm-granite/granite-4.0-tiny-preview"
-
-# Load model and tokenizer
-tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
-model = AutoModelForCausalLM.from_pretrained(MODEL_PATH, torch_dtype=torch.float32)
-
-def generate_text(prompt, max_tokens=200):
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+from transformers import AutoProcessor, AutoTokenizer, AutoModelForCausalLM
+from langgraph.graph import StateGraph
+# --- Model and Workflow Setup ---
+model_id = "ibm-granite/granite-4.0-tiny-preview"
+processor = AutoProcessor.from_pretrained(model_id)
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(model_id)
+def generate_with_granite(prompt: str, max_tokens: int = 200, use_gpu: bool = False) -> str:
+    device = torch.device("cuda" if use_gpu and torch.cuda.is_available() else "cpu")
     model.to(device)
-    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+    messages = [{"role": "user", "content": prompt}]
+    inputs = processor.apply_chat_template(
+        messages,
+        add_generation_prompt=True,
+        tokenize=True,
+        return_tensors="pt"
+    ).to(device)
     outputs = model.generate(
-        **inputs,
+        input_ids=inputs,
         max_new_tokens=max_tokens,
         do_sample=True,
         temperature=0.7,
         top_p=0.9,
         pad_token_id=tokenizer.eos_token_id
     )
-    return tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
-
-def select_genre_and_tone(user_input):
-    prompt = f"""You are a creative assistant. The user wants to write a short animated story.
+    generated = tokenizer.decode(outputs[0][inputs.shape[-1]:], skip_special_tokens=True)
+    return generated.strip()
+def select_genre_node(state: dict) -> dict:
+    prompt = f"""
+You are a creative assistant. The user wants to write a short animated story.
 Based on the following input, suggest a suitable genre and tone for the story.
-User Input: {user_input}
+User Input: {state['user_input']}
 Respond in this format:
 Genre: <genre>
 Tone: <tone>
-"""
-    response = generate_text(prompt, max_tokens=64)
+""".strip()
+    response = generate_with_granite(prompt)
     genre, tone = None, None
     for line in response.splitlines():
         if "Genre:" in line:
             genre = line.split("Genre:")[1].strip()
         elif "Tone:" in line:
             tone = line.split("Tone:")[1].strip()
-    return genre or "Unknown", tone or "Unknown"
-
-def generate_outline(user_input, genre, tone):
-    prompt = f"""You are a creative writing assistant helping to write a short animated screenplay.
+    state["genre"] = genre
+    state["tone"] = tone
+    return state
+def generate_outline_node(state: dict) -> dict:
+    prompt = f"""
+You are a creative writing assistant helping to write a short animated screenplay.
 The user wants to write a story with the following details:
-Genre: {genre}
-Tone: {tone}
-Idea: {user_input}
+Genre: {state.get('genre')}
+Tone: {state.get('tone')}
+Idea: {state.get('user_input')}
 Write a brief plot outline (3–5 sentences) for the story.
-"""
-    return generate_text(prompt, max_tokens=128)
-
-def generate_scene(genre, tone, outline):
-    prompt = f"""You are a screenwriter.
+""".strip()
+    response = generate_with_granite(prompt, max_tokens=250)
+    state["outline"] = response
+    return state
+def generate_scene_node(state: dict) -> dict:
+    prompt = f"""
+You are a screenwriter.
 Based on the following plot outline, write a key scene from the story.
 Focus on a turning point or climax moment. Make the scene vivid, descriptive, and suitable for an animated short film.
-Genre: {genre}
-Tone: {tone}
-Outline: {outline}
+Genre: {state.get('genre')}
+Tone: {state.get('tone')}
+Outline: {state.get('outline')}
 Write the scene in prose format (not screenplay format).
-"""
-    return generate_text(prompt, max_tokens=128)
-
-def generate_dialogue(scene):
-    prompt = f"""You are a dialogue writer for an animated screenplay.
+""".strip()
+    response = generate_with_granite(prompt, max_tokens=300)
+    state["scene"] = response
+    return state
+def write_dialogue_node(state: dict) -> dict:
+    prompt = f"""
+You are a dialogue writer for an animated screenplay.
 Below is a scene from the story:
-{scene}
+{state.get('scene')}
 Write the dialogue between the characters in screenplay format.
 Keep it short, expressive, and suitable for a short animated film.
 Use character names (you may invent them if needed), and format as:
@@ -74,32 +86,49 @@ CHARACTER:
 Dialogue line
 CHARACTER:
 Dialogue line
-"""
-    return generate_text(prompt, max_tokens=128)
-
-# Flask app
+""".strip()
+    response = generate_with_granite(prompt, max_tokens=100)
+    state["dialogue"] = response
+    return state
+def with_progress(fn, label, index, total):
+    def wrapper(state):
+        print(f"\n[{index}/{total}] Starting: {label}")
+        start = time.time()
+        result = fn(state)
+        duration = time.time() - start
+        print(f"[{index}/{total}] Completed: {label} in {duration:.2f} seconds")
+        return result
+    return wrapper
+def build_workflow():
+    graph = StateGraph(dict)
+    graph.add_node("select_genre", with_progress(select_genre_node, "Select Genre", 1, 4))
+    graph.add_node("generate_outline", with_progress(generate_outline_node, "Generate Outline", 2, 4))
+    graph.add_node("generate_scene", with_progress(generate_scene_node, "Generate Scene", 3, 4))
+    graph.add_node("write_dialogue", with_progress(write_dialogue_node, "Write Dialogue", 4, 4))
+    graph.set_entry_point("select_genre")
+    graph.add_edge("select_genre", "generate_outline")
+    graph.add_edge("generate_outline", "generate_scene")
+    graph.add_edge("generate_scene", "write_dialogue")
+    graph.set_finish_point("write_dialogue")
+    return graph.compile()
+workflow = build_workflow()
+# --- Flask App ---
 app = Flask(__name__)
 CORS(app)
-
 @app.route("/generate-story", methods=["POST"])
 def generate_story():
     data = request.get_json()
-    user_input = data.get("user_input", "")
+    user_input = data.get("user_input")
     if not user_input:
-        return jsonify({"error": "Missing user_input"}), 400
-
-    genre, tone = select_genre_and_tone(user_input)
-    outline = generate_outline(user_input, genre, tone)
-    scene = generate_scene(genre, tone, outline)
-    dialogue = generate_dialogue(scene)
-
+        return jsonify({"error": "Missing 'user_input' in request."}), 400
+    initial_state = {"user_input": user_input}
+    final_state = workflow.invoke(initial_state)
     return jsonify({
-        "genre": genre,
-        "tone": tone,
-        "outline": outline,
-        "scene": scene,
-        "dialogue": dialogue
+        "genre": final_state.get("genre"),
+        "tone": final_state.get("tone"),
+        "outline": final_state.get("outline"),
+        "scene": final_state.get("scene"),
+        "dialogue": final_state.get("dialogue")
     })
-
 if __name__ == "__main__":
-    app.run(host="0.0.0.0", port=8000, debug=True)
+    app.run(host="0.0.0.0", port=int(os.environ.get("PORT", 8000)), debug=True)