Picture.py - master - github_better/GenAI

Picture.py
326 строк12 KB
github_better
New
22 дек 2025, 08:52
22 дек 2025, 08:52b68d8c4
import replicate
from dotenv import load_dotenv
import os
from pathlib import Path
import uuid
import logging

from Lib.Debugger import debug_print
from Lib.Translate_text import translate_text

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

current_file = Path(__file__)
project_root = current_file.parent.parent
env_path = project_root / "Secret.env"
load_dotenv(env_path)

TOKEN = os.getenv("PICTURE_TOKEN")
if not TOKEN:
    debug_print(logger, "ERROR here NOT Token")

client = replicate.Client(api_token=TOKEN)

# Создаём папку Image_buffer, если её нет
image_buffer_path = current_file.parent / "Image_buffer"
image_buffer_path.mkdir(parents=True, exist_ok=True)

# --- Стили генерации картинки (их всего 2) ---
STYLES = {
    1: "photorealistic, highly detailed, professional photography, 8k resolution",
    2: "cartoon style, stylized illustration, bold colors, simplified shapes",
}

# --- Стиль → логический ключ модели (оба стиля идут в одну и ту же модель) ---
STYLE_TO_MODEL = {
    1: "flux_1_1_pro_ultra",  # фотореализм
    2: "flux_1_1_pro_ultra",  # cartoon
}

# --- Конфиги моделей Replicate ---
MODEL_CONFIGS = {
    "flux_1_1_pro_ultra": {
        "endpoint": "black-forest-labs/flux-1.1-pro-ultra",
        "defaults": {
            "num_inference_steps": 30,
            "guidance_scale": 4.5,
            "output_format": "png",
        },
    },
}

# --- Негативный промпт для генерации ---
NEGATIVE_PROMPT_BASE = (
    "low quality, blurry, out of focus, noisy, grainy, lowres, pixelated, mosaic, "
    "distorted, deformed, disfigured, mutated, bad anatomy, bad proportions, "
    "extra limbs, extra arms, extra legs, extra heads, extra hands, extra fingers, "
    "missing fingers, poorly drawn face, poorly drawn hands, ugly face, cloned face, "
    "duplicate people, crowd, group of people, many people, background people, random people, "
    "child, children, baby, doll, mannequin, "
    "watermark, logo, text, caption, subtitle, border, frame, cropped, out of frame"
)

def generate_picture(prompt: str, style_id: int = 1) -> dict:
    debug_print(logger)
    """
    Генерация изображения через Replicate с учётом стиля (цифрой).

    :param prompt: итоговый промпт для картинки (английский)
    :param style_id: ID стиля (1 или 6)
    :return: dict с информацией о сгенерированном файле
    """

    # Определяем модель по стилю (если передали что-то не то — падаем в стиль 1)
    if style_id not in STYLES:
        style_id = 1

    model_key = STYLE_TO_MODEL.get(style_id, "flux_1_1_pro_ultra")
    cfg = MODEL_CONFIGS.get(model_key)
    if not cfg:
        raise ValueError(f"Unknown model_key for style_id={style_id}: {model_key}")

    # Базовые параметры модели + промпт
    input_data = dict(cfg["defaults"])
    input_data["prompt"] = prompt
    input_data["negative_prompt"] = NEGATIVE_PROMPT_BASE

    # Вызов Replicate
    debug_print(logger, f"Running model {cfg['endpoint']} for style_id={style_id}")
    output = client.run(cfg["endpoint"], input=input_data)

    # Сохранение результата
    name = str(uuid.uuid4())
    file_path = image_buffer_path / f"{name}.png"

    # Предполагаем байтовый поток; при другом формате адаптируй сохранение
    with open(file_path, "wb") as f:
        f.write(output.read())

    debug_print(logger, f"Image saved to {file_path}")
    return {
        "path": str(file_path),
        "name": name,
        "model_key": model_key,
        "style_id": style_id,
    }


def delete_image(name_im: str) -> bool:
    debug_print(logger)
    """
    Удаляет изображение из папки Image_buffer по имени файла.
    
    :param name_im: имя файла (с расширением или без)
    :return: True если удаление успешно, False если файл не найден или ошибка
    """

    try:
        # Если передан полный путь, извлекаем только имя файла
        if "/" in name_im or "\\" in name_im:
            name_im = Path(name_im).name

        # Если расширение не указано, добавляем .png
        if not name_im.endswith((".webp", ".png", ".jpg", ".jpeg")):
            name_im = f"{name_im}.png"

        file_path = image_buffer_path / name_im

        # Проверяем существование файла
        if not file_path.exists():
            debug_print(logger, f"Файл не найден: {file_path}")
            return False

        # Защита от path traversal
        if not str(file_path.resolve()).startswith(str(image_buffer_path.resolve())):
            debug_print(logger, f"Попытка удаления файла вне Image_buffer: {file_path}")
            return False

        file_path.unlink()
        debug_print(logger, f"✅ Изображение успешно удалено: {file_path}")
        return True

    except Exception as e:
        debug_print(logger, f"❌ Ошибка при удалении изображения {name_im}: {e}")
        return False


def cleanup_old_images(days_old: int = 7) -> int:
    debug_print(logger, f"Запуск очистки изображений старше {days_old} дней")
    """
    Удаляет изображения старше указанного количества дней из Image_buffer.
    
    :param days_old: количество дней, после которых файлы считаются устаревшими
    :return: количество удаленных файлов
    """

    import time

    deleted_count = 0
    current_time = time.time()
    max_age_seconds = days_old * 24 * 60 * 60

    try:
        # При желании можешь заменить на "*.png", если используешь только PNG
        for file_path in image_buffer_path.glob("*.png"):
            if not file_path.is_file():
                continue

            file_age = current_time - file_path.stat().st_mtime
            if file_age > max_age_seconds:
                try:
                    file_path.unlink()
                    deleted_count += 1
                    debug_print(logger, f"Удален старый файл: {file_path.name}")
                except Exception as e:
                    debug_print(logger, f"Ошибка удаления {file_path.name}: {e}")

        debug_print(logger, f"✅ Очистка завершена. Удалено файлов: {deleted_count}")
        return deleted_count

    except Exception as e:
        debug_print(logger, f"❌ Ошибка при очистке старых изображений: {e}")
        return deleted_count


def choose_style(val: int) -> str:
    debug_print(logger)
    """
    Возвращает описание стиля генерации изображений по номеру.
    Args:
        val: номер стиля (1 или 2)
    Returns:
        Описание стиля для передачи в генератор изображений
        
    :param индекс параметра
    """

    return STYLES.get(val, STYLES[1])  # По умолчанию фотореализм


def pretty_style(lang: str) -> str:
    debug_print(logger)
    """
    Возвращает красиво отформатированный список доступных стилей генерации
    (только 1 и 2).
    
    :param сырой текст стилей
    """

    # Идём по реальным ключам STYLES — сейчас это только 1 и 6
    style_ids = sorted(STYLES.keys())

    names_en = [
        "Photorealistic",
        "Cartoon Style",
    ]
    descriptions_en = [
        "Highly realistic photography",
        "Stylized animation",
    ]
    emojis = ["🎥", "🎨"]

    # Перевод
    translated_names = [translate_text(name, "en", lang) for name in names_en]
    translated_descriptions = [translate_text(desc, "en", lang) for desc in descriptions_en]

    style_names = {}
    descriptions = {}
    for i, style_id in enumerate(style_ids):
        style_names[style_id] = f"{emojis[i]} {translated_names[i]}"
        descriptions[style_id] = translated_descriptions[i]

    header_en = "🎨 Available styles:\n\n"
    header = translate_text(header_en, "en", lang)

    result = header
    for style_id in style_ids:
        result += f"{style_id}. {style_names[style_id]}\n {descriptions[style_id]}\n\n"

    return result.strip()


def main_picture_prompt(
    current_location: str,
    current_time: str,
    last_picture_context: str,
    generation_style: str,
    description_of_scene: str = "",
    character_state: str = "",
) -> str:
    debug_print(logger)
    """
    Генерирует промпт для создания RPG иллюстрации.
    ...
    character_state: краткое описание состояния героя (статы + инвентарь), low‑priority контекст
    
    :param current_location: текущая локация 
    :param current_time: текущее время
    :param last_picture_context: контекст последней, сгенерированной картинке
    :param generation_style: стиль генерации
    :param description_of_scene: описание сцены
    :param character_state: Инвентарь+статы перса
    :return: итоговый промпт
    """

    try:
        prompt = f"""VISUAL STYLE:
{generation_style}

TASK:
Create a single detailed RPG illustration for this scene.

MAIN SUBJECTS:
- Clearly show the main character and any other important people in the scene.
- Faces and bodies must be clearly visible and anatomically correct.
- Do not add new characters that are not implied by the context.

SCENE SETTING:
- Location: {current_location}
- In‑world time: {current_time}
"""

        if description_of_scene and description_of_scene.strip():
            prompt += f"""

CURRENT ACTION TO ILLUSTRATE:
{description_of_scene.strip()}
"""

        # Лёгкий контекст по статам и инвентарю (без сильного акцента)
        if character_state and character_state.strip():
            prompt += f"""

CHARACTER STATE (OPTIONAL CONTEXT):
{character_state.strip()}
"""

        prompt += """
COMPOSITION GUIDELINES:
- Cinematic, immersive framing with a clear focal point on the main character or key action.
- Use lighting and atmosphere that match the time and mood of the scene.
- Include environmental details that reflect the character's journey and current situation.
- Avoid extreme close‑ups unless clearly implied by the action.

FINAL REQUIREMENTS:
- High quality, detailed rendering
- Natural perspective and realistic proportions
- Rich colors and atmospheric effects
- No text, captions or watermarks in the image
"""

        if last_picture_context and last_picture_context.strip():
            prompt += f"""

VISUAL CONTINUITY:
- Keep character appearance consistent with previous images:
{last_picture_context}
- Preserve overall visual tone and world style, adapt only the pose, angle and situation.
"""

        debug_print(logger, f"Main picture prompt generated with style {generation_style!r}")
        return prompt

    except Exception as e:
        debug_print(logger, f"ERROR in main_picture_prompt: {e}")
        return "ERROR"