replacebg / scene_planner.py
Munaf1987's picture
Update scene_planner.py
4bbe186 verified
raw
history blame
699 Bytes
# In scene_planner.py
from transformers import pipeline
scene_splitter = pipeline("text-generation", model="microsoft/phi-2")
def plan_scenes(script):
prompt = (
"Split this story into 10-15 cartoon scenes. Each should have a JSON with 'prompt' (image description) "
"and 'dialogue' (line of narration). Make it suitable for animation.\nStory:\n" + script
)
response = scene_splitter(prompt, max_new_tokens=1024, do_sample=False)[0]['generated_text']
try:
import json
scenes = json.loads(response)
except:
scenes = [{"prompt": line.strip(), "dialogue": line.strip()} for line in script.split('.') if line.strip()]
return scenes[:15]