main.py

import base64
import os
import time
import uuid
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from gradio_client import Client, handle_file

app = FastAPI(title="DailyDrop ML Proxy")

# Глобальная переменная, но БЕЗ мгновенной инициализации
hf_client = None

class TryOnRequest(BaseModel):
    user_image_b64: str
    garment_image_b64: str
    description: str = "A stylish garment"

class TryOnResponse(BaseModel):
    status: str
    result_image_base64: str
    processing_time_sec: float

def save_b64_to_temp(b64_str: str) -> str:
    if "," in b64_str:
        b64_str = b64_str.split(",")[1]
    file_data = base64.b64decode(b64_str)
    file_path = f"/tmp/{uuid.uuid4().hex}.jpg"
    with open(file_path, "wb") as f:
        f.write(file_data)
    return file_path

def image_to_base64(image_path: str) -> str:
    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode('utf-8')

def get_hf_client():
    """Ленивая (Lazy) инициализация клиента с поддержкой токена"""
    global hf_client
    if hf_client is None:
        print("⏳ [ML Proxy] Первое подключение к кластеру IDM-VTON...")
        hf_token = os.environ.get("HF_TOKEN") # Берем токен из секретов, если он есть
        hf_client = Client("yisol/IDM-VTON", hf_token=hf_token)
    return hf_client

@app.post("/api/v1/generate", response_model=TryOnResponse)
async def generate_look(request: TryOnRequest):
    print("📥 [ML Proxy] Задача получена (Base64)")
    start_time = time.time()
    
    # Сохраняем присланные картинки во временные файлы
    user_path = save_b64_to_temp(request.user_image_b64)
    garment_path = save_b64_to_temp(request.garment_image_b64)

    try:
        client = get_hf_client()
        
        result_paths = client.predict(
            dict({"background": handle_file(user_path), "layers": [], "composite": None}),
            handle_file(garment_path),
            request.description,
            True,   # auto-mask
            True,   # auto-crop
            30,     # denoise steps
            42,     # seed
            api_name="/tryon"
        )
        
        generated_temp_path = result_paths[0]
        img_base64 = image_to_base64(generated_temp_path)
        
        for p in [user_path, garment_path, generated_temp_path]:
            if os.path.exists(p): os.remove(p)
                
        proc_time = round(time.time() - start_time, 2)
        print(f"✅ [ML Proxy] Успех! Заняло {proc_time} сек.")

        return TryOnResponse(
            status="success",
            result_image_base64=img_base64,
            processing_time_sec=proc_time
        )

    except Exception as e:
        print(f"⚠️ [ML Proxy] Внешнее API перегружено или недоступно: {str(e)}")
        print("🔄 [ML Proxy] Включаем режим Fallback (возвращаем исходное фото для PoC)")
        
        # ЕСЛИ ПУБЛИЧНАЯ СЕТЬ УПАЛА, МЫ НЕ РОНЯЕМ ПРИЛОЖЕНИЕ!
        # Возвращаем просто фотографию человека, чтобы Бэкенд и Мобилка продолжили работать
        proc_time = round(time.time() - start_time, 2)
        return TryOnResponse(
            status="fallback",
            result_image_base64=request.user_image_b64, # Отдаем то, что прислали
            processing_time_sec=proc_time
        )