Rizz-Therapy / helpers.py
inventwithdean's picture
add costume selection
a4f1c5e
Raw
History Blame Contribute Delete
3.22 kB
import re
import io
import zipfile
import json
def parse_dialogue_tags(dialogue):
"""
Parses the dialogues with <look_at:character> tags
and returns clean dialogue and tags
"""
parts = re.split(r'(<look_at:[a-zA-Z]+>)', dialogue)
current_word_idx = 0
clean_text = ""
tags = []
for part in parts:
if part.startswith("<look_at:"):
# It is a look_at tag
character = part[9:-1].lower()
tags.append({
"character": character.lower(),
"word_idx": current_word_idx
})
else:
# It is normal part of sentence.
words = part.strip().split()
current_word_idx += len(words)
clean_text += part + " "
return clean_text.strip(), tags
def create_episode_zip(episode_payload, audio_chunks):
"""
Args:
episode_payload: Json containing the dialogues
audio_chunks: dict containing file names mapped to np.float32 arrays
"""
zip_buffer = io.BytesIO()
with zipfile.ZipFile(zip_buffer, "a", zipfile.ZIP_DEFLATED, False) as zip_file:
zip_file.writestr("episode.json", json.dumps(episode_payload))
for filename, audio_array in audio_chunks.items():
zip_file.writestr(filename, audio_array.tobytes())
return zip_buffer.getvalue()
def postprocess_script(script, character_names: dict[str, str]):
"""
Replaces all named characters with their roles, so unity can use them.
Args:
script: The parsed JSON script
character_names: A dict like {"therapist": "emily", "male": "mike", "female": "mia"}
"""
# Now it's {"emily": "therapist", "mike": "male", "mia": "female"}
name_to_role = {name.lower(): role for role, name in character_names.items()}
for dialogue in script.get("dialogues", []):
char_name = dialogue.get("character", "").lower()
if char_name in name_to_role:
dialogue["character"] = name_to_role[char_name]
for look_target in dialogue.get("lookAt", []):
target_name = look_target.get("character", "").lower()
if target_name in name_to_role:
look_target["character"] = name_to_role[target_name]
return script
expected_schema = {
"type": "object",
"properties": {
"dialogues": {
"type": "array",
"items": {
"type": "object",
"properties": {
"character": {"type": "string"},
"dialogue": {"type": "string"},
"expression": {"type": "string"}
},
"required": ["character", "dialogue", "expression"]
}
}
},
"required": ["dialogues"]
}
therapistCostumes = [
"therapist_beauty",
"therapist_suit",
"therapist_futuristic",
"therapist_vampire",
"therapist_huggingface",
]
maleCostumes = [
"male_charisma",
"male_suit",
"male_shirt",
"male_vampire",
"male_hoodie",
]
femaleCostumes = [
"female_charisma",
"female_suit",
"female_futuristic",
"female_vampire",
"female_offgrid",
]