add logs

app.py

@@ -1,6 +1,7 @@
 import os
 import io
 import time
+from datetime import datetime
 import spaces
 import torch
 import gradio as gr
@@ -14,35 +15,45 @@ from fastapi import FastAPI, Response
 import uvicorn
 from pydantic import BaseModel
 
+def print_with_time(*args, **kwargs):
+    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}{message}", **kwargs)
+
+print_with_time("Старт сервиса")
+
 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()}")
+print_with_time("="*50)
+print_with_time(f"PyTorch version: {torch.__version__}")
+print_with_time(f"CUDA available: {torch.cuda.is_available()}")
+print_with_time(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)
+    print_with_time(f"Current device: {torch.cuda.current_device()}")
+    print_with_time(f"Device name: {torch.cuda.get_device_name(0)}")
+print_with_time("="*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...")
+        print_with_time(f"Используется VRAM: {torch.cuda.memory_allocated() / 1024 ** 3:.2f} GB")
+        print_with_time(f"Доступно VRAM: {torch.cuda.memory_reserved() / 1024 ** 3:.2f} GB")
+        print_with_time(f"Пиковое использование VRAM: {torch.cuda.max_memory_allocated() / 1024 ** 3:.2f} GB")
+        print_with_time(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")
+        print_with_time(f"Очистка кеша CUDA завершена.")
+        print_with_time(f"Используется VRAM: {torch.cuda.memory_allocated() / 1024 ** 3:.2f} GB")
+        print_with_time(f"Доступно VRAM: {torch.cuda.memory_reserved() / 1024 ** 3:.2f} GB")
+        print_with_time(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")
+print_with_time(f"Загрузка модели FLUX.1-dev")
 pipe = FluxPipeline.from_pretrained(
             # pretrained_model_name_or_path = local_path,
             "black-forest-labs/FLUX.1-dev",
@@ -53,69 +64,71 @@ pipe = FluxPipeline.from_pretrained(
             # variant="fp16",
             use_safetensors=True
         )
-print(f"Загрузка модели FLUX.1-dev завершена за {time.time() - start_time:.2f} секунд")
+print_with_time(f"Загрузка модели FLUX.1-dev завершена за {time.time() - start_time:.2f} секунд")
 
 start_time = time.time()
-print(f"Загрузка LoRA")
+print_with_time(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} секунд")
+print_with_time(f"Загрузка LoRA завершена за {time.time() - start_time:.2f} секунд")
 
 pipe.fuse_lora(lora_scale=0.5)
 pipe.to(device)
 
 start_time = time.time()
-print(f"Загрузка модели fal/AuraSR-v2")
+print_with_time(f"Загрузка модели fal/AuraSR-v2")
 aura_sr = AuraSR.from_pretrained("fal/AuraSR-v2")
-print(f"Загрузка модели fal/AuraSR-v2 завершена за {time.time() - start_time:.2f} секунд")
+print_with_time(f"Загрузка модели fal/AuraSR-v2 завершена за {time.time() - start_time:.2f} секунд")
 
 start_time = time.time()
-print(f"Загрузка модели briaai/RMBG-1.4")
+print_with_time(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_with_time(f"Загрузка модели briaai/RMBG-1.4 завершена за {time.time() - start_time:.2f} секунд")
 
-print(f"Загрузка всех моделей завершена за {time.time() - start_full_time:.2f} секунд")
+print_with_time(f"Загрузка всех моделей завершена за {time.time() - start_full_time:.2f} секунд")
 
 @spaces.GPU()
 def generate_image(object_name, remove_bg=True, upscale=True):
     try:
-        # Формирование промпта
+        print_with_time("Формирование промпта")
         object_name = translate_ru_en(object_name)
         prompt = create_template_prompt(object_name)
         
         # Для имитации генерации (можно заменить на реальный вызов ComfyUI API)
-        print(f"Генерация иконки для объекта: {object_name}")
-        print(f"Промпт: {prompt[:100]}...")
+        print_with_time(f"Генерация иконки для объекта: {object_name}")
+        print_with_time(f"Промпт: {prompt[:100]}...")
         # print(f"Параметры: seed={seed}, steps={steps}, размер={width}x{height}")
-        print(f"Опции: remove_bg={remove_bg}")
+        print_with_time(f"Опции: remove_bg={remove_bg}")
         steps = int(os.getenv('STEPS')) if os.getenv('STEPS') is not None else 10
-        print(f"Шаги: {steps}")
+        print_with_time(f"Шаги: {steps}")
 
         clear_cuda()
 
         start_time = time.time()
-        print(f"Старт генерации изображения...")
+        print_with_time("Старт генерации изображения...")
         image = generate_image_stable(prompt, steps)
-        print(f"Генерация завершена за {time.time() - start_time:.2f} секунд")
+        print_with_time(f"Генерация завершена за {time.time() - start_time:.2f} секунд")
 
         if upscale :
             clear_cuda()
             start_time = time.time()
-            print(f"Старт апскейлера...")
+            print_with_time(f"Старт апскейлера...")
             image = aura_sr.upscale_4x_overlapped(image)
-            print(f"Апскейлинг завершен за {time.time() - start_time:.2f} секунд")
+            print_with_time(f"Апскейлинг завершен за {time.time() - start_time:.2f} секунд")
 
         if remove_bg :
             clear_cuda()
             start_time = time.time()
-            print(f"Старт удаления фона...")
+            print_with_time(f"Старт удаления фона...")
             image = bg_remover(image)
-            print(f"Фон удален за {time.time() - start_time:.2f} секунд")
+            print_with_time(f"Фон удален за {time.time() - start_time:.2f} секунд")
 
+        print_with_time("Генерация успешно завершена.")
         clear_cuda()
+        print_with_time("Отдаем готовое изображение")
         return image
     
     except Exception as e:
-        print(f"Ошибка при генерации изображения: {e}")
+        print_with_time(f"Ошибка при генерации изображения: {e}")
         traceback.print_exc()
         clear_cuda()
         return None
@@ -147,7 +160,7 @@ def translate_ru_en(text: str):
         return translator.translate(text)
 
     except Exception as e:
-        print(f"Ошибка перевода: {e}")
+        print_with_time(f"Ошибка перевода: {e}")
         traceback.print_exc()
         return text
 
@@ -155,7 +168,6 @@ def load_text(file_name):
     with open(file_name, 'r', encoding='utf-8') as f:
         return f.read()
 
-
 fastapi_app = FastAPI()
 
 # Модель данных для POST-запроса

requirements.txt

@@ -21,3 +21,4 @@ deep_translator
 fastapi
 uvicorn
 pydantic
+pytz