Spaces:
Runtime error
Runtime error
File size: 2,837 Bytes
55b3b1b |
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 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 |
from __future__ import annotations
import os
from typing import List, Optional, Tuple
from .models import Storyboard
from .placeholders import synthesize_character_card
DEFAULT_IMAGE_MODEL = os.environ.get("CINEGEN_CHARACTER_MODEL", "gemini-2.5-flash-image")
def _load_google_client(api_key: Optional[str]):
if not api_key:
return None
try:
from google import genai
return genai.Client(api_key=api_key)
except Exception: # pragma: no cover - optional dependency
return None
class CharacterDesigner:
def __init__(self, api_key: Optional[str] = None):
self.api_key = api_key or os.environ.get("GOOGLE_API_KEY")
self.client = _load_google_client(self.api_key)
def design(self, storyboard: Storyboard) -> Tuple[List[Tuple[str, str]], Storyboard]:
gallery: List[Tuple[str, str]] = []
for character in storyboard.characters:
gallery.append(self._refresh_reference(character, storyboard.style))
return gallery, storyboard
def redesign_character(self, storyboard: Storyboard, character_id: str) -> Tuple[Tuple[str, str], Storyboard]:
target = next((char for char in storyboard.characters if char.identifier == character_id), None)
if not target:
raise ValueError(f"Character {character_id} not found.")
card = self._refresh_reference(target, storyboard.style)
return card, storyboard
def _refresh_reference(self, character, style: str) -> Tuple[str, str]:
image_path = None
if self.client:
image_path = self._try_generate(character, style)
if not image_path:
image_path = synthesize_character_card(character, style)
character.reference_image = image_path
caption = f"{character.name} — {character.role}"
return image_path, caption
def _try_generate(self, character, style: str) -> Optional[str]: # pragma: no cover
prompt = (
f"Create a portrait for {character.name}, a {character.role} in a {style} short film. "
f"Traits: {', '.join(character.traits)}. Description: {character.description}."
)
try:
response = self.client.models.generate_content(
model=DEFAULT_IMAGE_MODEL,
contents=[prompt],
)
for part in response.parts:
if getattr(part, "inline_data", None):
image = part.as_image()
tmp_dir = os.path.join("/tmp", "cinegen-characters")
os.makedirs(tmp_dir, exist_ok=True)
path = os.path.join(tmp_dir, f"{character.identifier.lower()}.png")
image.save(path)
return path
except Exception:
return None
return None
|