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