initial commit

README.md

@@ -1,5 +1,5 @@
 ---
-title: IconDDDzilla Bot Backend
+title: IconDDDzilla_Bot_Backend
 emoji: 🚀
 colorFrom: blue
 colorTo: indigo

app.py

@@ -0,0 +1,259 @@
+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()

prompt.txt

@@ -0,0 +1 @@
+3D {object_name} icon: Minimal geo, brushed aluminum, matte acrylic, rubber, 3-point, 8k, PBR, isometric

requirements.txt

@@ -0,0 +1,21 @@
+--extra-index-url https://download.pytorch.org/whl/cu114
+torch
+torchvision
+gradio==3.50.2
+pillow>=9.5.0
+numpy>=1.24.0
+huggingface_hub>=0.17.0
+safetensors>=0.3.0
+accelerate>=0.18.0
+requests>=2.28.0
+opencv-python>=4.6.0
+transformers
+diffusers
+peft
+safetensors
+sentencepiece
+scikit-image
+rembg
+aura-sr
+deep_translator
+pytz

space.yaml

@@ -0,0 +1,8 @@
+title: IconDDDzilla_Bot_Backend
+emoji: 🚀
+colorFrom: blue
+colorTo: indigo
+sdk: gradio
+sdk_version: 3.50.2
+app_file: app.py
+pinned: false

style.css

@@ -0,0 +1,187 @@
+:root {
+    --primary: #5A54F9;
+    --primary-dark: #4D47D1;
+    --secondary: #8B84FF;
+    --text: #111827;
+    --light: #F9FAFB;
+}
+
+.gradio-container {
+    background: linear-gradient(135deg, #F9FAFB 0%, #E5E7EB 100%) !important;
+    font-family: 'Segoe UI', system-ui, -apple-system, sans-serif !important;
+    max-width: 800px !important;
+    margin: 20px auto !important;
+    border-radius: 16px !important;
+    box-shadow: 0 10px 25px -5px rgba(90, 84, 249, 0.1) !important;
+}
+
+.dark .gradio-container {
+    background: linear-gradient(135deg, #1F2937 0%, #111827 100%) !important;
+    --text: #F9FAFB;
+    --light: #1F2937;
+}
+
+.gr-block {
+    border: none !important;
+    background: white !important;
+    border-radius: 12px !important;
+    box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.05) !important;
+    padding: 20px !important;
+    margin-bottom: 24px !important;
+    transition: all 0.3s ease !important;
+}
+
+.dark .gr-block {
+    background: rgba(30, 41, 59, 0.8) !important;
+    backdrop-filter: blur(10px) !important;
+}
+
+.gr-block:hover {
+    transform: translateY(-2px) !important;
+    box-shadow: 0 10px 15px -3px rgba(0, 0, 0, 0.1) !important;
+}
+
+.gr-button {
+    background: linear-gradient(135deg, var(--primary) 0%, var(--primary-dark) 100%) !important;
+    color: white !important;
+    border: none !important;
+    border-radius: 8px !important;
+    padding: 12px 24px !important;
+    font-weight: 600 !important;
+    text-transform: uppercase !important;
+    letter-spacing: 0.5px !important;
+    transition: all 0.3s ease !important;
+    box-shadow: 0 4px 6px -1px rgba(90, 84, 249, 0.3) !important;
+}
+
+.gr-button:hover {
+    transform: translateY(-2px) !important;
+    box-shadow: 0 10px 15px -3px rgba(90, 84, 249, 0.4) !important;
+}
+
+.gr-button:active {
+    transform: scale(0.98) !important;
+}
+
+.gr-textbox, .gr-number, .gr-slider, .gr-dropdown {
+    border: 1px solid #e2e8f0 !important;
+    border-radius: 8px !important;
+    padding: 12px 16px !important;
+    background: var(--light) !important;
+    transition: all 0.3s ease !important;
+}
+
+.gr-textbox:focus, .gr-number:focus, .gr-dropdown:focus {
+    border-color: var(--primary) !important;
+    box-shadow: 0 0 0 2px rgba(90, 84, 249, 0.2) !important;
+    outline: none !important;
+}
+
+h1, h2, h3, h4 {
+    color: var(--text) !important;
+    font-weight: 600 !important;
+    margin-top: 0 !important;
+}
+
+.gr-markdown {
+    color: var(--text) !important;
+    line-height: 1.6 !important;
+}
+
+.progress-bar {
+    height: 6px !important;
+    border-radius: 3px !important;
+    background: linear-gradient(90deg, var(--primary), #93c5fd) !important;
+}
+
+.tabs {
+    background: transparent !important;
+}
+
+.tab-button {
+    border-radius: 8px 8px 0 0 !important;
+    margin-right: 4px !important;
+    transition: all 0.3s ease !important;
+}
+
+.tab-button.selected {
+    background: white !important;
+    color: var(--primary) !important;
+    font-weight: 600 !important;
+}
+
+.dark .tab-button.selected {
+    background: rgba(30, 41, 59, 0.8) !important;
+}
+
+.gr-checkbox {
+    margin: 12px 0 !important;
+}
+
+.gr-checkbox label {
+    display: flex !important;
+    align-items: center !important;
+    cursor: pointer !important;
+    position: relative !important;
+    padding-left: 35px !important;
+    min-height: 25px !important;
+}
+
+.gr-checkbox input[type="checkbox"] {
+    position: absolute !important;
+    opacity: 0 !important;
+    cursor: pointer !important;
+    height: 0 !important;
+    width: 0 !important;
+}
+
+.checkmark {
+    position: absolute !important;
+    top: 0 !important;
+    left: 0 !important;
+    height: 24px !important;
+    width: 24px !important;
+    background-color: var(--light) !important;
+    border: 2px solid var(--primary) !important;
+    border-radius: 6px !important;
+    transition: all 0.3s ease !important;
+}
+
+.gr-checkbox:hover .checkmark {
+    transform: scale(1.05) !important;
+    box-shadow: 0 0 0 3px rgba(90, 84, 249, 0.1) !important;
+}
+
+.gr-checkbox input:checked ~ .checkmark {
+    background: linear-gradient(135deg, var(--primary) 0%, var(--primary-dark) 100%) !important;
+    border-color: var(--primary-dark) !important;
+}
+
+.checkmark:after {
+    content: "" !important;
+    position: absolute !important;
+    display: none !important;
+}
+
+.gr-checkbox input:checked ~ .checkmark:after {
+    display: block !important;
+}
+
+.gr-checkbox .checkmark:after {
+    left: 8px !important;
+    top: 4px !important;
+    width: 6px !important;
+    height: 12px !important;
+    border: solid white !important;
+    border-width: 0 2px 2px 0 !important;
+    transform: rotate(45deg) !important;
+}
+
+.dark .gr-checkbox .checkmark {
+    background-color: rgba(30, 41, 59, 0.5) !important;
+    border-color: #93c5fd !important;
+}
+
+.dark .gr-checkbox input:checked ~ .checkmark {
+    background: linear-gradient(135deg, var(--primary) 0%, #1e8a50 100%) !important;
+}