164 lines
7.2 KiB
Python
164 lines
7.2 KiB
Python
# Copyright (C) 2025 AIDC-AI
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
"""
|
|
Character Memory component for web UI
|
|
"""
|
|
|
|
import streamlit as st
|
|
from typing import List
|
|
|
|
from web.i18n import tr
|
|
from web.utils.streamlit_helpers import safe_rerun
|
|
from pixelle_video.services.quality import CharacterMemory, Character, CharacterType
|
|
|
|
|
|
def render_active_characters_badge():
|
|
"""Render a small badge showing active characters count (for Home page)"""
|
|
if "character_memory" not in st.session_state:
|
|
return
|
|
|
|
memory = st.session_state.character_memory
|
|
active_count = sum(1 for c in memory.characters if getattr(c, 'is_active', True))
|
|
|
|
if active_count > 0:
|
|
st.info(f"👤 {tr('character.registered_count', count=active_count)} (Active)", icon="✨")
|
|
|
|
|
|
def render_character_memory():
|
|
"""Render the character memory management interface"""
|
|
st.markdown(f"### {tr('character.title')}")
|
|
st.caption(tr("character.desc"))
|
|
|
|
# Initialize CharacterMemory in session state if not present
|
|
if "character_memory" not in st.session_state:
|
|
st.session_state.character_memory = CharacterMemory()
|
|
|
|
memory = st.session_state.character_memory
|
|
|
|
# --- Character Registration Form ---
|
|
with st.expander(tr("character.add_button"), expanded=False):
|
|
with st.form("add_character_form", clear_on_submit=True):
|
|
col1, col2 = st.columns(2)
|
|
|
|
with col1:
|
|
name = st.text_input(tr("character.name"), key="new_char_name")
|
|
char_type = st.selectbox(
|
|
tr("character.type"),
|
|
options=[t.value for t in CharacterType],
|
|
format_func=lambda x: tr(f"character.type.{x}"),
|
|
key="new_char_type"
|
|
)
|
|
|
|
with col2:
|
|
features = st.text_input(tr("character.features"), placeholder="e.g., glasses, scar, red hat", key="new_char_features")
|
|
# Image upload placeholder (Streamlit doesn't support persistent file paths easily in memory objects)
|
|
ref_image = st.file_uploader(tr("character.ref_image"), type=["png", "jpg", "jpeg"], key="new_char_image")
|
|
|
|
appearance = st.text_area(tr("character.appearance"), placeholder="e.g., young man with short black hair", key="new_char_appearance")
|
|
clothing = st.text_area(tr("character.clothing"), placeholder="e.g., blue t-shirt and jeans", key="new_char_clothing")
|
|
|
|
submit = st.form_submit_button(tr("character.add_button"), use_container_width=True)
|
|
|
|
if submit:
|
|
if not name:
|
|
st.error(tr("error.missing_field", field=tr("character.name")))
|
|
else:
|
|
import os
|
|
import shutil
|
|
from pathlib import Path
|
|
|
|
feature_list = [f.strip() for f in features.split(",") if f.strip()]
|
|
|
|
# Handle ref image upload
|
|
ref_path = None
|
|
if ref_image:
|
|
# Save to a temporary storage or artifacts dir
|
|
upload_dir = Path("/Users/yuanjiantsui/projects/Pixelle-Video/temp/characters")
|
|
upload_dir.mkdir(parents=True, exist_ok=True)
|
|
ref_path = str(upload_dir / f"{name}_{ref_image.name}")
|
|
with open(ref_path, "wb") as f:
|
|
f.write(ref_image.getbuffer())
|
|
|
|
char = memory.register_character(
|
|
name=name,
|
|
appearance_description=appearance,
|
|
clothing_description=clothing,
|
|
distinctive_features=feature_list,
|
|
character_type=CharacterType(char_type)
|
|
)
|
|
|
|
if ref_path:
|
|
char.add_reference_image(ref_path, set_as_primary=True)
|
|
|
|
st.success(f"✅ {name} registered!")
|
|
safe_rerun()
|
|
|
|
# --- Character List ---
|
|
st.markdown("---")
|
|
characters = memory.characters
|
|
|
|
if not characters:
|
|
st.info(tr("character.no_characters"))
|
|
else:
|
|
st.markdown(tr("character.registered_count", count=len(characters)))
|
|
|
|
# Display characters in a grid
|
|
cols = st.columns(3)
|
|
for i, char in enumerate(characters):
|
|
with cols[i % 3]:
|
|
with st.container(border=True):
|
|
# Header with Status Toggle
|
|
head_col, action_col = st.columns([3, 1])
|
|
with head_col:
|
|
st.markdown(f"**{char.name}**")
|
|
st.caption(tr(f"character.type.{char.character_type.value}"))
|
|
with action_col:
|
|
# Toggle for activation (defensive access for cached objects)
|
|
current_active = getattr(char, 'is_active', True)
|
|
new_state = st.toggle("On", value=current_active, key=f"tog_{char.id}", label_visibility="collapsed")
|
|
if new_state != current_active:
|
|
char.is_active = new_state
|
|
safe_rerun()
|
|
|
|
# Image Preview
|
|
if char.primary_reference:
|
|
st.image(char.primary_reference, use_container_width=True)
|
|
else:
|
|
st.empty()
|
|
|
|
if char.appearance_description:
|
|
st.caption(f"📝 {char.appearance_description}")
|
|
|
|
if char.clothing_description:
|
|
st.caption(f"👕 {char.clothing_description}")
|
|
|
|
if char.distinctive_features:
|
|
st.markdown(", ".join([f"`{f}`" for f in char.distinctive_features]))
|
|
|
|
st.markdown("---")
|
|
# Delete Button at bottom
|
|
if st.button("🗑️ " + tr("history.task_card.delete"), key=f"del_{char.id}", use_container_width=True):
|
|
# Basic delete logic
|
|
del memory._characters[char.id]
|
|
if char.name.lower() in memory._name_index:
|
|
del memory._name_index[char.name.lower()]
|
|
safe_rerun()
|
|
|
|
# Provide consistency prompt suggestion if characters exist
|
|
if characters:
|
|
with st.expander("💡 Character Consistency Prompt", expanded=False):
|
|
st.code(
|
|
", ".join([char.get_prompt_injection() for char in characters]),
|
|
language="text"
|
|
)
|