app.py

import sqlite3
import streamlit as st
import random
from typing import List, Dict, Any

DB_PATH = "fine_prompt_sdxl.db"


def get_db_connection() -> sqlite3.Connection:
    """Abre una conexión a la base de datos SQLite."""
    conn = sqlite3.connect(DB_PATH)
    conn.row_factory = sqlite3.Row
    return conn


def load_franchises() -> List[sqlite3.Row]:
    """Carga todas las franquicias de la base de datos."""
    with get_db_connection() as conn:
        return conn.execute("SELECT * FROM franchises").fetchall()


def load_characters(franchise_id: int) -> List[sqlite3.Row]:
    """Carga los personajes asociados a una franquicia."""
    with get_db_connection() as conn:
        return conn.execute(
            "SELECT * FROM characters WHERE franchise_id = ?", (franchise_id,)
        ).fetchall()


def load_categories(tables: List[str]) -> Dict[str, List[str]]:
    """Carga las categorías de las tablas indicadas."""
    categories = {}
    with get_db_connection() as conn:
        for table in tables:
            rows = conn.execute(f"SELECT value FROM {table}").fetchall()
            categories[table] = [row["value"] for row in rows]
    return categories


def load_special_values(special_tables: List[str]) -> Dict[str, List[str]]:
    """Carga los valores especiales de las tablas indicadas."""
    special_values = {}
    with get_db_connection() as conn:
        for table in special_tables:
            rows = conn.execute(f"SELECT value FROM {table}").fetchall()
            special_values[table] = [row["value"] for row in rows]
    return special_values


def sanitize_character_name(name: str) -> str:
    """Elimina la parte 'from [franquicia]' del nombre del personaje solo para la lista."""
    return name.split(" from ")[0].strip() if " from " in name else name


def apply_weight(value: str) -> str:
    """Agrega un peso aleatorio entre 0.8 y 1.5 a un elemento."""
    weight = round(random.uniform(0.8, 1.5), 1)
    return f"({value}:{weight})"


def generate_prompt(
    character: Dict[str, Any],
    categories: Dict[str, List[str]],
    special_selection: str,
    special_values: Dict[str, List[str]],
) -> str:
    """Genera un prompt combinando el personaje, las categorías y los valores especiales."""
    elements = {
        "hairstyle": apply_weight(random.choice(categories["hairstyle"])),
        "outfit": apply_weight(random.choice(categories["outfit"])),
        "scenario": apply_weight(random.choice(categories["scenario"])),
        "emotion": apply_weight(random.choice(categories["emotion"])),
        "pose": apply_weight(random.choice(categories["pose"])),
        "extras": ", ".join(
            apply_weight(extra) for extra in random.sample(categories["extras"], k=2)
        ),
        "lighting": apply_weight(random.choice(categories["lighting"])),
        "distance": apply_weight(random.choice(categories["distance"])),
        "angle": apply_weight(random.choice(categories["angle"])),
        "special_elements": apply_weight(random.choice(categories["special_elements"])),
        "style": apply_weight(random.choice(categories["style"])),
    }

    special_section = ""
    if special_selection in ("special_all", "special_s_e", "both"):
        special_section = "(nsfw:1.5), "
        if special_selection == "both":
            special_all_values = random.sample(special_values.get("special_all", []), 2)
            special_se_values = random.sample(special_values.get("special_s_e", []), 1)
            all_special_values = special_all_values + special_se_values
            special_section += (
                ", ".join(apply_weight(value) for value in all_special_values) + ", "
            )  # esto sirve para
        else:
            values = special_values.get(
                "special_all" if special_selection == "special_all" else "special_s_e",
                [],
            )
            selected_values = random.sample(values, 2)
            special_section += (
                ", ".join(apply_weight(value) for value in selected_values) + ", "
            )

    core_tags = character["core_tags"] if "core_tags" in character.keys() else ""
    core_tags = random.choice(
        [tag.strip() for tag in core_tags.split(";") if tag.strip()]
    )
    core_tags_section = f"{core_tags}, " if core_tags else ""
    prompt = (
        f"{character['name']}, {core_tags_section}"
        f"{elements['hairstyle']}, wearing {elements['outfit']}, positioned in a {elements['scenario']}, "
        f"striking a pose of {elements['pose']}, feeling {elements['emotion']}. "
        f"Scene elements include {elements['extras']}, {special_section}"
        f"{elements['angle']}, {elements['distance']}, {elements['lighting']}, "
        f"{elements['special_elements']}, "
        "masterpiece, high score, great score, absurdres, best quality, highres"
    )
    return prompt


st.title("Prompt Generator")

franchises = load_franchises()
franchise_name_to_id = {franchise["name"]: franchise["id"] for franchise in franchises}
selected_franchise = st.selectbox(
    "Select a franchise:", list(franchise_name_to_id.keys())
)
selected_franchise_id = franchise_name_to_id[selected_franchise]

characters = load_characters(selected_franchise_id)

character_name_to_data = {
    sanitize_character_name(character["name"]): character for character in characters
}

selected_character_name = st.selectbox(
    "Select a character:", list(character_name_to_data.keys())
)

selected_character = character_name_to_data[selected_character_name]


status_color_map = {
    "recognized": "green",
    "partial": "yellow",
    "not recognized": "red",
}
character_status = (
    selected_character["status"] if "status" in selected_character.keys() else "unknown"
)
status_color = status_color_map.get(character_status, "gray")
st.markdown(
    f"<p style='color: {status_color}; font-size: 16px;'>Status: {character_status.capitalize()}</p>",
    unsafe_allow_html=True,
)

special_selection = st.radio(
    "Include special categories:", ["none", "special_all", "special_s_e", "both"]
)

if st.button("Generate Prompt"):
    tables = [
        "hairstyle",
        "outfit",
        "scenario",
        "emotion",
        "pose",
        "extras",
        "lighting",
        "distance",
        "angle",
        "special_elements",
        "style",
    ]
    categories = load_categories(tables)
    special_values = {}
    if special_selection in ("special_all", "special_s_e", "both"):
        special_values = load_special_values(["special_all", "special_s_e"])
    prompt = generate_prompt(
        selected_character, categories, special_selection, special_values
    )
    st.code(prompt, language="python", wrap_lines=True)