(self, messages, **kwargs)
| 98 | ] |
| 99 | |
| 100 | async def run(self, messages, **kwargs): |
| 101 | logger.info('Segmenting script to sentences.') |
| 102 | if os.path.exists(os.path.join(self.work_dir, 'segments.txt')): |
| 103 | return messages |
| 104 | with open(os.path.join(self.work_dir, 'script.txt'), 'r') as f: |
| 105 | script = f.read() |
| 106 | with open(os.path.join(self.work_dir, 'topic.txt'), 'r') as f: |
| 107 | topic = f.read() |
| 108 | |
| 109 | image_prompt = '' |
| 110 | if self.config.background != 'image': |
| 111 | image_prompt = f'\n\n背景图片无需生成,是纯色:{self.config.background}\n\n' |
| 112 | |
| 113 | query = (f'原始主题:\n\n{topic}\n\n' |
| 114 | f'原始脚本:\n\n{script}\n\n' |
| 115 | f'{image_prompt}' |
| 116 | f'请完成你的动画分镜设计:\n') |
| 117 | messages = await super().run(query, **kwargs) |
| 118 | response = messages[-1].content |
| 119 | if '```json' in response: |
| 120 | response = response.split('```json')[1].split('```')[0] |
| 121 | elif '```' in response: |
| 122 | response = response.split('```')[1].split('```')[0] |
| 123 | segments = json.loads(response) |
| 124 | if self.config.foreground == 'image': |
| 125 | segments = await self.add_images(segments, topic, script, **kwargs) |
| 126 | |
| 127 | for i, segment in enumerate(segments): |
| 128 | assert 'content' in segment |
| 129 | if self.config.background == 'image': |
| 130 | assert 'background' in segment or 'video' in segment |
| 131 | else: |
| 132 | segment['background'] = self.config.background |
| 133 | if 'video' in segment: |
| 134 | segment.pop('background', None) |
| 135 | segment.pop('manim', None) |
| 136 | segment.pop(self.engine, None) |
| 137 | segment.pop('foreground', None) |
| 138 | logger.info( |
| 139 | f'\n场景 {i}\n' |
| 140 | f'内容:{segment["content"]}\n' |
| 141 | f'图片要求:{segment.get("background", "无背景")}\n' |
| 142 | f'视频要求:{segment.get("video", "非视频片段")}\n' |
| 143 | f'动画要求:{segment.get(self.engine, segment.get("manim", "无动画"))}' |
| 144 | ) |
| 145 | with open(os.path.join(self.work_dir, 'segments.txt'), 'w') as f: |
| 146 | f.write(json.dumps(segments, indent=4, ensure_ascii=False)) |
| 147 | return messages |
| 148 | |
| 149 | async def add_images(self, segments, topic, script, **kwargs): |
| 150 |
no test coverage detected