Louis Labeyrie
e03daea1f3
app.py relied on parsing stdout for the save path printed by the pipeline. After PR #21 replaced that print with logger.info(), the stdout parsing returned None, causing "Aucune image générée détectée". Now check the known save path directly (APP_DIR / save_path) and only fall back to stdout parsing if the file isn't found.
145 lines
4.9 KiB
Python
145 lines
4.9 KiB
Python
import streamlit as st
|
|
import subprocess
|
|
import sys
|
|
import shlex
|
|
from pathlib import Path
|
|
from PIL import Image
|
|
|
|
APP_DIR = Path(__file__).resolve().parent
|
|
PIPELINE_SCRIPT = APP_DIR / "prompt_to_card_pipeline.py"
|
|
TEXT_CLEANER_PATH = APP_DIR / "text-cleaner" / "text_cleaning_pipeline.py"
|
|
INFER_SCRIPT_PATH = APP_DIR / "clean-text-to-keywords" / "infer_json_usage.py"
|
|
CHECKPOINT_PATH = APP_DIR / "pokemon_card_lora" / "training_history.pt"
|
|
TEMPLATE_PATH = APP_DIR / "clean-text-to-keywords" / "json_template_example.json"
|
|
IMAGE_EXTENSIONS = (".png", ".jpg", ".jpeg", ".webp", ".bmp")
|
|
|
|
|
|
def _extract_image_from_stdout(stdout: str) -> Path | None:
|
|
for line in reversed(stdout.splitlines()):
|
|
# Try the whole line, then the part after the last colon
|
|
# (handles "Card generated and saved to: generated_card.png")
|
|
raw = line.strip().strip("\"'")
|
|
candidates = [raw]
|
|
if ":" in raw:
|
|
candidates.append(raw.rsplit(":", 1)[1].strip().strip("\"'"))
|
|
|
|
for text in candidates:
|
|
if not text:
|
|
continue
|
|
candidate = Path(text)
|
|
if not candidate.is_absolute():
|
|
candidate = APP_DIR / candidate
|
|
if candidate.suffix.lower() in IMAGE_EXTENSIONS and candidate.exists():
|
|
return candidate
|
|
|
|
return None
|
|
|
|
|
|
|
|
def run_prompt_pipeline(prompt_text: str) -> tuple[Path | None, str, list[str]]:
|
|
save_path = "generated_card.png"
|
|
cmd = [
|
|
sys.executable, "prompt_to_card_pipeline.py",
|
|
prompt_text,
|
|
"--text-cleaner-path", "text-cleaner/text_cleaning_pipeline.py",
|
|
"--infer-script-path", "clean-text-to-keywords/infer_json_usage.py",
|
|
"--checkpoint", "pokemon_card_lora",
|
|
"--template", "clean-text-to-keywords/json_template_example.json",
|
|
"--generator-module", "card_generator_adapter.py",
|
|
"--device", "cuda",
|
|
"--save-path", save_path,
|
|
"--print-json",
|
|
]
|
|
|
|
result = subprocess.run(
|
|
cmd,
|
|
cwd=APP_DIR,
|
|
capture_output=True,
|
|
text=True,
|
|
check=False,
|
|
)
|
|
|
|
full_output = (result.stdout or "") + ("\n" + result.stderr if result.stderr else "")
|
|
|
|
if result.returncode != 0:
|
|
return None, full_output.strip() or "Erreur inconnue pendant le pipeline.", cmd
|
|
|
|
image_path = APP_DIR / save_path
|
|
if not image_path.exists():
|
|
image_path = _extract_image_from_stdout(result.stdout or "")
|
|
return image_path, full_output.strip(), cmd
|
|
|
|
# ------------------------------------------------------------------ #
|
|
# Configuration #
|
|
# ------------------------------------------------------------------ #
|
|
|
|
st.set_page_config(
|
|
page_title="Générateur de Carte Pokémon",
|
|
page_icon=Image.open(Path(__file__).with_name("pokeball.png")),
|
|
layout="centered",
|
|
)
|
|
|
|
logo_col, title_col = st.columns([1, 6], vertical_alignment="center")
|
|
|
|
with logo_col:
|
|
st.image(Image.open(Path(__file__).with_name("pokeball.png")), width=72)
|
|
|
|
with title_col:
|
|
st.title("Générateur de Carte Pokémon")
|
|
st.markdown("Décrivez votre Pokémon en langage naturel et laissez la magie opérer !")
|
|
|
|
# ------------------------------------------------------------------ #
|
|
# Saisie utilisateur #
|
|
# ------------------------------------------------------------------ #
|
|
|
|
raw_text = st.text_area(
|
|
label="Description de votre Pokémon",
|
|
placeholder=(
|
|
"Ex: My pokemon is called Pyrokar! Its a huge fire dragon with massive "
|
|
"red wings and shoots flames from its mouth... super fast n aggressive >:("
|
|
),
|
|
height=180,
|
|
)
|
|
|
|
st.markdown(
|
|
"""
|
|
<style>
|
|
div.stButton > button {
|
|
background-color: #d62828;
|
|
color: white;
|
|
border: 1px solid #b91c1c;
|
|
}
|
|
div.stButton > button:hover {
|
|
background-color: #b91c1c;
|
|
border-color: #991b1b;
|
|
color: white;
|
|
}
|
|
</style>
|
|
""",
|
|
unsafe_allow_html=True,
|
|
)
|
|
|
|
generate = st.button(" Générer la carte", use_container_width=True)
|
|
|
|
# ------------------------------------------------------------------ #
|
|
# Pipeline #
|
|
# ------------------------------------------------------------------ #
|
|
|
|
if generate:
|
|
if not raw_text.strip():
|
|
st.warning("Veuillez entrer une description avant de générer.")
|
|
else:
|
|
with st.spinner("Génération de la carte Pokémon..."):
|
|
image, logs, _cmd = run_prompt_pipeline(raw_text)
|
|
|
|
if image is not None:
|
|
st.image(image, caption="Carte Pokémon générée", width="stretch")
|
|
if logs:
|
|
with st.expander("Logs pipeline"):
|
|
st.code(logs)
|
|
else:
|
|
st.error("Aucune image générée détectée. Vérifiez les chemins du pipeline.")
|
|
if logs:
|
|
with st.expander("Logs pipeline"):
|
|
st.code(logs)
|