| 1 | import logging |
| 2 | import os |
| 3 | import asyncio |
| 4 | from langchain_core.prompts import ChatPromptTemplate |
| 5 | from langchain_core.output_parsers import PydanticOutputParser |
| 6 | from langchain.chat_models.base import BaseChatModel |
| 7 | from langchain.chat_models import init_chat_model |
| 8 | from pydantic import BaseModel, Field |
| 9 | from typing import List, Optional, Dict |
| 10 | from interfaces import CharacterInScene, ImageOutput |
| 11 | from langchain_core.messages import HumanMessage, SystemMessage |
| 12 | |
| 13 | |
| 14 | |
| 15 | prompt_template_front = \ |
| 16 | """ |
| 17 | Generate a full-body, front-view portrait of character {identifier} based on the following description, with a pure white background. Use a wide 16:9 landscape canvas, not a vertical portrait canvas. The character should be centered in the image, occupying the middle of the wide frame with enough horizontal empty space. Gazing straight ahead. Standing with arms relaxed at sides. Natural expression. |
| 18 | Features: {features} |
| 19 | Style: {style} |
| 20 | """ |
| 21 | |
| 22 | prompt_template_side = \ |
| 23 | """ |
| 24 | Generate a full-body, side-view portrait of character {identifier} based on the provided front-view portrait, with a pure white background. Use a wide 16:9 landscape canvas, not a vertical portrait canvas. The character should be centered in the image, occupying the middle of the wide frame with enough horizontal empty space. Facing left. Standing with arms relaxed at sides. |
| 25 | """ |
| 26 | |
| 27 | prompt_template_back = \ |
| 28 | """ |
| 29 | Generate a full-body, back-view portrait of character {identifier} based on the provided front-view portrait, with a pure white background. Use a wide 16:9 landscape canvas, not a vertical portrait canvas. The character should be centered in the image, occupying the middle of the wide frame with enough horizontal empty space. No facial features should be visible. |
| 30 | """ |
| 31 | |
| 32 | |
| 33 | class CharacterPortraitsGenerator: |
| 34 | def __init__( |
| 35 | self, |
| 36 | image_generator, |
| 37 | ): |
| 38 | self.image_generator = image_generator |
| 39 | |
| 40 | |
| 41 | async def generate_front_portrait( |
| 42 | self, |
| 43 | character: CharacterInScene, |
| 44 | style: str, |
| 45 | ) -> ImageOutput: |
| 46 | features = "(static) " + (character.static_features or "") + "; (dynamic) " + (character.dynamic_features or "") |
| 47 | prompt = prompt_template_front.format( |
| 48 | identifier=character.identifier_in_scene, |
| 49 | features=features, |
| 50 | style=style, |
| 51 | ) |
| 52 | image_output = await self.image_generator.generate_single_image( |
| 53 | prompt=prompt, |
| 54 | # size="512x512", |
| 55 | ) |
| 56 | return image_output |
| 57 | |
| 58 | async def generate_side_portrait( |
| 59 | self, |
| 60 | character: CharacterInScene, |
| 61 | front_image_path: str, |
| 62 | ) -> ImageOutput: |
| 63 | prompt = prompt_template_side.format( |
| 64 | identifier=character.identifier_in_scene, |
| 65 | ) |
| 66 | image_output = await self.image_generator.generate_single_image( |
| 67 | prompt=prompt, |
| 68 | reference_image_paths=[front_image_path], |
| 69 | # size="1024x1024", |
| 70 | ) |
| 71 | return image_output |
| 72 | |
| 73 | |
| 74 | async def generate_back_portrait( |
| 75 | self, |
| 76 | character: CharacterInScene, |
| 77 | front_image_path: str, |
| 78 | ) -> ImageOutput: |
| 79 | prompt = prompt_template_back.format( |
| 80 | identifier=character.identifier_in_scene, |
| 81 | ) |
| 82 | image_output = await self.image_generator.generate_single_image( |
| 83 | prompt=prompt, |
| 84 | reference_image_paths=[front_image_path], |
| 85 | # size="512x512", |
| 86 | ) |
| 87 | return image_output |