app.py

import os
import time
from datetime import datetime
import pytz
import spaces
import torch
import gradio as gr
import traceback
import threading
from transformers import pipeline
from huggingface_hub import login
from diffusers import FluxPipeline
from aura_sr import AuraSR
from deep_translator import GoogleTranslator

def log(*args, **kwargs):
    tid = threading.get_ident()
    tz = pytz.timezone(kwargs.pop('timezone', 'Europe/Moscow'))
    time_fmt = kwargs.pop('time_format', '%Y-%m-%d %H:%M:%S')
    separator = kwargs.pop('sep', ' - ')
    moscow_time = datetime.now(tz).strftime(time_fmt)
    message = separator.join(str(arg) for arg in args)
    print(f"[{moscow_time}]{separator}[tid={tid}]{separator}{message}", **kwargs)

log("Старт сервиса")

login(token=os.getenv('HF_TOKEN'))

log("=" * 50)
log(f"PyTorch version: {torch.__version__}")
log(f"CUDA available: {torch.cuda.is_available()}")
log(f"GPU count: {torch.cuda.device_count()}")
if torch.cuda.is_available():
    log(f"Current device: {torch.cuda.current_device()}")
    log(f"Device name: {torch.cuda.get_device_name(0)}")
log("=" * 50)

def clear_cuda():
    if torch.cuda.is_available():
        log(f"Используется VRAM: {torch.cuda.memory_allocated() / 1024 ** 3:.2f} GB")
        log(f"Доступно VRAM: {torch.cuda.memory_reserved() / 1024 ** 3:.2f} GB")
        log(f"Пиковое использование VRAM: {torch.cuda.max_memory_allocated() / 1024 ** 3:.2f} GB")
        log(f"Очистка кеша CUDA...")
        torch.cuda.empty_cache()
        torch.cuda.reset_peak_memory_stats()
        log(f"Очистка кеша CUDA завершена.")
        log(f"Используется VRAM: {torch.cuda.memory_allocated() / 1024 ** 3:.2f} GB")
        log(f"Доступно VRAM: {torch.cuda.memory_reserved() / 1024 ** 3:.2f} GB")
        log(f"Пиковое использование VRAM: {torch.cuda.max_memory_allocated() / 1024 ** 3:.2f} GB")

device = "cuda" if torch.cuda.is_available() else "cpu"

start_full_time = time.time()
start_time = time.time()
log(f"Загрузка модели FLUX.1-dev")
pipe = FluxPipeline.from_pretrained(
    # pretrained_model_name_or_path = local_path,
    "black-forest-labs/FLUX.1-dev",
    torch_dtype=torch.bfloat16,  # Используем bfloat16 для A100
    # low_cpu_mem_usage=True,  # Экономия памяти
    # device_map="balanced",
    # local_files_only=True
    # variant="fp16",
    use_safetensors=True
)
log(f"Загрузка модели FLUX.1-dev завершена за {time.time() - start_time:.2f} секунд")

start_time = time.time()
log(f"Загрузка LoRA")
pipe.load_lora_weights("Shakker-Labs/FLUX.1-dev-LoRA-add-details", weight_name="FLUX-dev-lora-add_details.safetensors")
log(f"Загрузка LoRA завершена за {time.time() - start_time:.2f} секунд")

pipe.fuse_lora(lora_scale=0.5)
pipe.to(device)

start_time = time.time()
log(f"Загрузка модели fal/AuraSR-v2")
aura_sr = AuraSR.from_pretrained("fal/AuraSR-v2")
log(f"Загрузка модели fal/AuraSR-v2 завершена за {time.time() - start_time:.2f} секунд")

start_time = time.time()
log(f"Загрузка модели briaai/RMBG-1.4")
bg_remover = pipeline("image-segmentation", "briaai/RMBG-1.4", trust_remote_code=True)
log(f"Загрузка модели briaai/RMBG-1.4 завершена за {time.time() - start_time:.2f} секунд")

log(f"Загрузка всех моделей завершена за {time.time() - start_full_time:.2f} секунд")

@spaces.GPU()
def generate_image(object_name, remove_bg=True, upscale=True):
    try:
        log("Формирование промпта")
        object_name = translate_ru_en(object_name)
        prompt = create_template_prompt(object_name)

        # Для имитации генерации (можно заменить на реальный вызов ComfyUI API)
        log(f"Генерация иконки для объекта: {object_name}")
        log(f"Промпт: {prompt[:100]}...")
        # print(f"Параметры: seed={seed}, steps={steps}, размер={width}x{height}")
        log(f"Опции: remove_bg={remove_bg}")
        steps = int(os.getenv('STEPS')) if os.getenv('STEPS') is not None else 10
        log(f"Шаги: {steps}")

        clear_cuda()

        start_time = time.time()
        log("Старт генерации изображения...")
        image = generate_image_stable(prompt, steps)
        log(f"Генерация завершена за {time.time() - start_time:.2f} секунд")

        if upscale:
            clear_cuda()
            start_time = time.time()
            log(f"Старт апскейлера...")
            image = aura_sr.upscale_4x_overlapped(image)
            log(f"Апскейлинг завершен за {time.time() - start_time:.2f} секунд")

        if remove_bg:
            clear_cuda()
            start_time = time.time()
            log(f"Старт удаления фона...")
            image = bg_remover(image)
            log(f"Фон удален за {time.time() - start_time:.2f} секунд")

        log(f"Генерация успешно завершена для запроса: {prompt}")
        clear_cuda()
        log(f"Отдаем готовое изображение для запроса: {prompt}")
        return image

    except Exception as e:
        log(f"Ошибка при генерации изображения: {e}")
        traceback.print_exc()
        clear_cuda()
        return None

def generate_image_stable(prompt, steps):
    return pipe(
        prompt,
        height=1024,
        width=1024,
        guidance_scale=3.5,
        num_inference_steps=steps,
        generator=torch.Generator(device).manual_seed(0),
        num_images_per_prompt=1
    ).images[0]

def create_template_prompt(object_name):
    template = load_text("prompt.txt")
    return template.format(object_name=object_name)

def translate_ru_en(text: str):
    try:
        # Проверка на кириллицу (если включено)
        if not any('\u0400' <= char <= '\u04FF' for char in text):
            return text

        # Создаем переводчик
        translator = GoogleTranslator(source="ru", target="en")
        # Выполняем перевод
        return translator.translate(text)

    except Exception as e:
        log(f"Ошибка перевода: {e}")
        traceback.print_exc()
        return text

def load_text(file_name):
    with open(file_name, 'r', encoding='utf-8') as f:
        return f.read()

custom_css = load_text("style.css")

# Создание интерфейса Gradio
with gr.Blocks(title="3D Icon Generator", css=custom_css, theme=gr.themes.Soft()) as app:
    gr.Markdown("# iconDDDzilla")
    gr.Markdown("### Create 3d icons with transparent background in one click!")

    with gr.Row():
        with gr.Column():
            # Входные параметры
            object_input = gr.Textbox(label="Object name",
                                      placeholder="Type object name (for example: calendar, phone, camera)")
            remove_bg_checkbox = gr.Checkbox(label="Remove background", value=True)
            upscale = gr.Checkbox(label="Upscale", value=True)

            # with gr.Accordion("Расширенные настройки", open=False):
            #     custom_prompt = gr.Textbox(label="Пользовательский промпт", placeholder="Оставьте пустым для использования шаблона", lines=3)
            #
            #     with gr.Row():
            #         seed = gr.Number(label="Seed", value=276789180904019, precision=0)
            #         steps = gr.Slider(minimum=1, maximum=5, value=5, step=1, label="Шаги")
            #
            #     with gr.Row():
            #         width = gr.Slider(minimum=512, maximum=2048, value=1024, step=64, label="Ширина")
            #         height = gr.Slider(minimum=512, maximum=2048, value=1024, step=64, label="Высота")
            #
            #     with gr.Row():
            #         #upscale_checkbox = gr.Checkbox(label="Применить апскейл", value=True)
            #         remove_bg_checkbox = gr.Checkbox(label="Удалить фон", value=False)

            # Кнопка генерации
            generate_btn = gr.Button("Run")

        with gr.Column():
            # Выходное изображение
            output_image = gr.Image(label="Image")

    # Примеры использования
    # examples = gr.Examples(
    #     examples=[
    #         ["calendar", "", 276789180904019, 5, 1024, 1024, True],
    #         ["camera", "", 391847529184, 5, 1024, 1024, True],
    #         ["smartphone", "", 654321987654, 5, 1024, 1024, True],
    #         ["headphones", "", 123456789012, 5, 1024, 1024, True],
    #     ],
    #     inputs=[
    #         object_input,
    #         custom_prompt,
    #         seed,
    #         steps,
    #         width,
    #         height,
    #         #upscale_checkbox,
    #         remove_bg_checkbox
    #     ],
    #     outputs=[output_image],
    #     fn=generate_image,
    # )

    # Информация о моделях
    # with gr.Accordion("Информация о используемых моделях", open=False):
    #     gr.Markdown("""
    #     ## Используемые модели
    #
    #     - **Основная модель:** [flux1-dev-fp8.safetensors](https://huggingface.co/lllyasviel/flux1_dev/blob/main/flux1-dev-fp8.safetensors) от Stability AI
    #     - **Модель апскейла:** [4x_NMKD-Superscale-SP_178000_G.pth](https://huggingface.co/gemasai/4x_NMKD-Superscale-SP_178000_G) для улучшения качества изображения
    #     - **Модель удаления фона:** [RMBG-1.4](https://huggingface.co/briaai/RMBG-1.4) для качественного удаления фона
    #
    #     Все модели автоматически загружаются при первом запуске приложения.
    #     """)

    # Привязка функции к нажатию кнопки
    generate_btn.click(
        fn=generate_image,
        inputs=[
            object_input,
            # custom_prompt,
            # seed,
            # steps,
            # width,
            # height,
            remove_bg_checkbox,
            upscale
        ],
        outputs=[output_image]
    )

# Запуск приложения
if __name__ == "__main__":
    app.launch()