app.py

import os
import time
import spaces
import torch
import gradio as gr
import traceback
from transformers import pipeline
from huggingface_hub import login
from diffusers import FluxPipeline
from aura_sr import AuraSR
from deep_translator import GoogleTranslator
from fastapi import FastAPI
import uvicorn

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

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

def clear_cuda():
    if torch.cuda.is_available():
        print(f"Используется VRAM: {torch.cuda.memory_allocated() / 1024 ** 3:.2f} GB")
        print(f"Доступно VRAM: {torch.cuda.memory_reserved() / 1024 ** 3:.2f} GB")
        print(f"Пиковое использование VRAM: {torch.cuda.max_memory_allocated() / 1024 ** 3:.2f} GB")
        print(f"Очистка кеша CUDA...")
        torch.cuda.empty_cache()
        torch.cuda.reset_peak_memory_stats()
        print(f"Очистка кеша CUDA завершена.")
        print(f"Используется VRAM: {torch.cuda.memory_allocated() / 1024 ** 3:.2f} GB")
        print(f"Доступно VRAM: {torch.cuda.memory_reserved() / 1024 ** 3:.2f} GB")
        print(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()
print(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
        )
print(f"Загрузка модели FLUX.1-dev завершена за {time.time() - start_time:.2f} секунд")

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

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

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

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

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

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

        clear_cuda()

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

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

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

        clear_cuda()
        return image
    
    except Exception as e:
        print(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(1),
        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:
        print(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()


fastapi_app = FastAPI()

@fastapi_app.get("/api/hello")
def get_users():
    return "Hello World!"


custom_css = load_text("style.css")

# Создание интерфейса Gradio
with gr.Blocks(title="3D Icon Generator", css=custom_css, theme=gr.themes.Default()) 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]
    )


main_app = gr.mount_gradio_app(fastapi_app, app, path="/")

# Запуск приложения
if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)




### OLD UNUSED CODE!!! BE CAREFUL!!! ###

# Создаем необходимые директории
# os.makedirs("models", exist_ok=True)
# os.makedirs("models/checkpoints", exist_ok=True)
# os.makedirs("models/loras", exist_ok=True)
# os.makedirs("models/upscale_models", exist_ok=True)
# os.makedirs("models/rembg", exist_ok=True)
# os.makedirs("outputs", exist_ok=True)
# os.makedirs("temp_uploads", exist_ok=True)


# Загрузка моделей с Hugging Face
# def download_model(repo_id, filename, local_dir):
#     """Загрузка модели с Hugging Face Hub"""
#     local_path = os.path.join(local_dir, filename)
#     if not os.path.exists(local_path):
#         print(f"Загрузка {filename} из {repo_id}...")
#         try:
#             file_path = hf_hub_download(
#                 repo_id=repo_id,
#                 filename=filename,
#                 local_dir=local_dir,
#                 local_dir_use_symlinks=False
#             )
#             print(f"Модель успешно загружена: {file_path}")
#             return file_path
#         except Exception as e:
#             print(f"Ошибка при загрузке модели {filename}: {e}")
#             return None
#     else:
#         print(f"Модель {filename} уже присутствует: {local_path}")
#         return local_path

# Вынесем загрузку моделей за пределы функции generate_image
# это критично для работы на ZeroGPU в Hugging Face Spaces
# def download_all_models():
#     models = {
#         "flux": download_model(
#             "lllyasviel/flux1_dev",
#             "flux1-dev-fp8.safetensors",
#             "models/checkpoints"
#         ),
#         # "upscale": download_model(
#         #     "gemasai/4x_NMKD-Superscale-SP_178000_G",
#         #     "4x_NMKD-Superscale-SP_178000_G.pth",
#         #     "models/upscale_models"
#         # ),
#         # "rembg": "briaai/RMBG-1.4"  # Используем модель RMBG-1.4 через transformers pipeline
#     }
#     return models

# def save_uploaded_file(file):
#     """Сохранение загруженного файла"""
#     if file is None:
#         return None
#
#     filename = os.path.join("temp_uploads", f"{random.randint(1000000, 9999999)}{os.path.splitext(file.name)[1]}")
#     with open(filename, "wb") as f:
#         f.write(file.read())
#     return filename

# Функция для получения значения по индексу (используется в ComfyUI)
# def get_value_at_index(obj, index):
#     """Получение значения из объекта по индексу из экспортированного ComfyUI кода"""
#     if isinstance(obj, list):
#         return obj[index]
#     elif isinstance(obj, tuple):
#         return obj[index]
#     elif isinstance(obj, dict):
#         return list(obj.values())[index]
#     else:
#         return obj[index]

# Загрузка моделей при запуске
#models = download_all_models()

# Инициализация pipeline для удаления фона
# try:
#     rembg_pipeline = pipeline("image-segmentation", model=models["rembg"], trust_remote_code=True)
#     print(f"Модель удаления фона успешно загружена: {models['rembg']}")
# except Exception as e:
#     print(f"Ошибка при загрузке модели удаления фона: {e}")
#     rembg_pipeline = None

# Укажите абсолютный путь к модели
# model_path = os.path.abspath("models/checkpoints/flux1-dev-fp8.safetensors")

# Проверьте существование ключевого файла
# if not os.path.exists(os.path.join(model_path, "model_index.json")):
#     raise FileNotFoundError(f"Модель не найдена по пути: {model_path}")

# local_path = os.path.join("models/checkpoints", "")