File size: 1,577 Bytes
d65b1bc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 |
import json
from pydantic import BaseModel
from typing import List
class Scene(BaseModel):
narration: str
image_prompts: List[str]
def num_images(self):
return len(self.image_prompts)
class Story(BaseModel):
scenes: List[Scene]
@classmethod
def from_dict(cls, data):
# Parse JSON string into list of dictionaries
json_data = data
for item in json_data:
if "image_prompts" in item and "narration" not in item:
for key in item.keys():
if key != "image_prompts":
item["narration"] = item.pop(key)
break
# Convert JSON data into list of Scene objects
scenes = [Scene(**item) for item in json_data]
# Return Story object with all scenes
return cls(scenes=scenes)
@classmethod
def parse_markdown(cls, markdown_str):
# Extract JSON part from markdown string
try:
json_str = markdown_str.split("```json")[1].split("```")[0].strip()
except:
pass
return cls.from_dict(data=json.loads(json_str))
@property
def all_narrations(self):
return [scene.narration for scene in self.scenes]
@property
def all_images(self):
# return [scene.images_prompt for scene in self.scenes]
results = []
for scene in self.scenes:
results.extend(scene.images_prompt)
return results
def num_images_in_scene(self, scene_index):
return self.scenes[scene_index].num_images()
|