Spaces:

nikitatupitsyn
/

iconDDDzilla_bot_backend

Running on Zero

File size: 10,797 Bytes

37193e0

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()