Update app.py

app.py

@@ -1,301 +1,269 @@
-# 🤖 HuggingFace FaceFusion API
-# Реальное API с максимальным качеством замены лиц (CPU)
+# 🤖 HuggingFace BFS Face Swap API (CPU Optimized)
 import io
 import os
 import logging
 import shutil
 from pathlib import Path
-import numpy as np
-from PIL import Image, ImageEnhance, ImageFilter
+import torch
+from PIL import Image
 import uvicorn
 from fastapi import FastAPI, File, UploadFile, Form, HTTPException
 from fastapi.responses import Response, JSONResponse
-
-# Попытка импорта FaceFusion с учётом возможных изменений в API
-FACEFUSION_AVAILABLE = False
-facefusion_import_error = None
-
-try:
-    # Пытаемся импортировать основные модули
-    from facefusion.face_analyser import get_many_faces
-    from facefusion.typing import Frame
-    
-    # Проверяем различные варианты импорта процессоров
-    try:
-        # Новый стиль (классы)
-        from facefusion.processors.frame.modules.face_swapper import FaceSwapper
-        from facefusion.processors.frame.modules.face_enhancer import FaceEnhancer
-        USE_CLASS_STYLE = True
-    except ImportError:
-        # Старый стиль (функции)
-        try:
-            from facefusion.processors.frame.modules.face_swapper import swap_face
-            from facefusion.processors.frame.modules.face_enhancer import enhance_face
-            USE_CLASS_STYLE = False
-        except ImportError:
-            # Падаем, если ничего не подошло
-            raise ImportError("Cannot import face_swapper or face_enhancer modules")
-    
-    FACEFUSION_AVAILABLE = True
-    print("✅ FaceFusion loaded successfully")
-except ImportError as e:
-    facefusion_import_error = str(e)
-    print(f"⚠️ FaceFusion not installed or incompatible: {e}")
+from diffusers import QwenImageEditPlusPipeline
+from huggingface_hub import snapshot_download
 
 # Настройка логирования
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-app = FastAPI(title="FaceFusion API", description="API для замены лиц с максимальным качеством", version="2.0.0")
+app = FastAPI(title="BFS Face Swap API", version="3.0.0")
 
-TEMP_DIR = Path("/tmp/facefusion")
+TEMP_DIR = Path("/tmp/bfs")
 TEMP_DIR.mkdir(exist_ok=True)
+MODEL_CACHE = Path("/app/model_cache")
+MODEL_CACHE.mkdir(exist_ok=True)
 
-# ----- Вспомогательные функции -----
-def enhance_image_quality(image_bytes: bytes) -> bytes:
-    """Дополнительное улучшение качества (постобработка)"""
-    try:
-        img = Image.open(io.BytesIO(image_bytes)).convert("RGB")
-        enhancer = ImageEnhance.Sharpness(img)
-        img = enhancer.enhance(1.2)
-        enhancer = ImageEnhance.Contrast(img)
-        img = enhancer.enhance(1.1)
-        img = img.filter(ImageFilter.UnsharpMask(radius=1, percent=120, threshold=3))
-        output = io.BytesIO()
-        img.save(output, format="JPEG", quality=98, optimize=True)
-        return output.getvalue()
-    except Exception as e:
-        logger.error(f"❌ Image enhancement error: {e}")
-        return image_bytes
+# Глобальные переменные
+pipe = None
+lora_loaded = False
 
-def save_upload_file(upload_file: UploadFile) -> Path:
-    """Сохраняет загруженный файл во временную папку"""
+# Конфигурация модели
+BASE_MODEL = "Qwen/Qwen-Image-Edit-2511"  # или 2509
+LORA_REPO = "Alissonerdx/BFS-Best-Face-Swap"
+LORA_FILE = "bfs_head_v5_2511_merged_version_rank_16_fp16.safetensors"  # Рекомендованная версия [citation:2][citation:4]
+
+def load_pipeline():
+    """Загрузка квантованной модели с CPU offload"""
+    global pipe, lora_loaded
+    
     try:
-        contents = upload_file.file.read()
-        unique_name = f"{os.urandom(8).hex()}_{upload_file.filename}"
-        file_path = TEMP_DIR / unique_name
-        with open(file_path, "wb") as f:
-            f.write(contents)
-        return file_path
+        logger.info("🔄 Loading 4-bit quantized model (first load takes 10-15 minutes)...")
+        
+        # Используем 4-bit версию для экономии RAM
+        pipe = QwenImageEditPlusPipeline.from_pretrained(
+            "toandev/Qwen-Image-Edit-2511-4bit",  # 4-bit quantized version
+            torch_dtype=torch.bfloat16,
+            low_cpu_mem_usage=True,
+            cache_dir=MODEL_CACHE
+        )
+        
+        # КРИТИЧЕСКИ ВАЖНО для CPU: включаем offload
+        pipe.enable_model_cpu_offload()
+        
+        # Загружаем BFS LoRA веса
+        logger.info("🔄 Loading BFS LoRA weights...")
+        
+        # Скачиваем LoRA если ещё нет
+        lora_path = MODEL_CACHE / LORA_FILE
+        if not lora_path.exists():
+            snapshot_download(
+                repo_id=LORA_REPO,
+                allow_patterns=[LORA_FILE],
+                local_dir=MODEL_CACHE
+            )
+        
+        pipe.load_lora_weights(str(lora_path))
+        lora_loaded = True
+        
+        logger.info("✅ Model and LoRA loaded successfully")
+        return True
+        
     except Exception as e:
-        logger.error(f"❌ File save error: {e}")
-        raise HTTPException(status_code=500, detail="Failed to save file")
-
-def read_image_numpy(path: Path) -> np.ndarray:
-    """Читает изображение как numpy array (RGB)"""
-    img = Image.open(path).convert("RGB")
-    return np.array(img)
-
-def write_numpy_image(arr: np.ndarray, quality: int = 98) -> bytes:
-    """Конвертирует numpy array в байты JPEG"""
-    img = Image.fromarray(arr.astype(np.uint8))
-    output = io.BytesIO()
-    img.save(output, format="JPEG", quality=quality, optimize=True)
-    return output.getvalue()
-
-# ----- Основные функции обработки через FaceFusion (адаптер) -----
-def process_face_swap(source_path: Path, target_path: Path,
-                      swapper_model: str = "inswapper_128",
-                      enhancer_model: str = None) -> np.ndarray:
-    """Замена лица с максимальным качеством"""
-    if not FACEFUSION_AVAILABLE:
-        logger.warning("FaceFusion not available, returning target with basic enhancements")
-        return read_image_numpy(target_path)
-
-    # Загружаем изображения
-    target_frame = read_image_numpy(target_path)
-    source_frame = read_image_numpy(source_path)
+        logger.error(f"❌ Failed to load model: {e}")
+        return False
 
-    # Детектируем лица
-    source_faces = get_many_faces(source_frame)
-    if not source_faces:
-        raise HTTPException(status_code=400, detail="No face found in source image")
-    source_face = source_faces[0]
-
-    target_faces = get_many_faces(target_frame)
-    if not target_faces:
-        raise HTTPException(status_code=400, detail="No face found in target image")
-
-    result_frame = target_frame.copy()
-
-    if USE_CLASS_STYLE:
-        # Используем классы FaceSwapper / FaceEnhancer
-        swapper = FaceSwapper(swapper_model)
-        for target_face in target_faces:
-            result_frame = swapper.process_frame(result_frame, source_face, target_face)
-        if enhancer_model:
-            enhancer = FaceEnhancer(enhancer_model)
-            result_frame = enhancer.process_frame(result_frame)
-    else:
-        # Старый стиль (функции)
-        for target_face in target_faces:
-            result_frame = swap_face(result_frame, source_face, target_face, model=swapper_model)
-        if enhancer_model:
-            result_frame = enhance_face(result_frame, model=enhancer_model)
-
-    return result_frame
-
-def process_face_enhance(image_path: Path, enhancer_model: str = "gfpgan_1.4") -> np.ndarray:
-    """Улучшение лица"""
-    if not FACEFUSION_AVAILABLE:
-        return read_image_numpy(image_path)
-    frame = read_image_numpy(image_path)
-    if USE_CLASS_STYLE:
-        enhancer = FaceEnhancer(enhancer_model)
-        return enhancer.process_frame(frame)
-    else:
-        return enhance_face(frame, model=enhancer_model)
-
-def process_face_analyse(image_path: Path) -> dict:
-    """Анализ лиц"""
-    if not FACEFUSION_AVAILABLE:
-        # Заглушка
-        img = Image.open(image_path)
-        w, h = img.size
-        return {
-            "faces_detected": 1,
-            "face_regions": [{"x": int(w*0.3), "y": int(h*0.1), "width": int(w*0.4), "height": int(h*0.4), "confidence": 0.95}],
-            "image_size": {"width": w, "height": h}
-        }
-    frame = read_image_numpy(image_path)
-    faces = get_many_faces(frame)
-    regions = []
-    for face in faces:
-        # Проверяем, какой формат у bounding box
-        if hasattr(face, 'bbox'):
-            bbox = face.bbox
-        elif hasattr(face, 'detection'):
-            bbox = face.detection
-        else:
-            # Пытаемся получить через стандартный способ
-            bbox = face.bbox if hasattr(face, 'bbox') else [0,0,0,0]
-        # Уверены, что bbox - список из 4 чисел
-        if len(bbox) == 4:
-            x1, y1, x2, y2 = map(int, bbox)
-            regions.append({
-                "x": x1,
-                "y": y1,
-                "width": x2 - x1,
-                "height": y2 - y1,
-                "confidence": getattr(face, 'det_score', 0.95)
-            })
-    h, w = frame.shape[:2]
-    return {
-        "faces_detected": len(faces),
-        "face_regions": regions,
-        "image_size": {"width": w, "height": h}
-    }
+def save_upload_file(upload_file: UploadFile) -> Path:
+    """Сохраняет загруженный файл"""
+    contents = upload_file.file.read()
+    unique_name = f"{os.urandom(8).hex()}_{upload_file.filename}"
+    file_path = TEMP_DIR / unique_name
+    with open(file_path, "wb") as f:
+        f.write(contents)
+    return file_path
 
-# ----- Эндпоинты -----
-@app.get("/")
-async def root():
-    return {"service": "FaceFusion API", "status": "running", "version": "2.0.0"}
+def optimize_image(image: Image.Image, max_size=1024) -> Image.Image:
+    """Оптимизация размера для ускорения"""
+    w, h = image.size
+    if max(w, h) > max_size:
+        scale = max_size / max(w, h)
+        new_w, new_h = int(w * scale), int(h * scale)
+        return image.resize((new_w, new_h), Image.Resampling.LANCZOS)
+    return image
 
-@app.get("/health")
-async def health():
-    return {"status": "ok", "facefusion_loaded": FACEFUSION_AVAILABLE}
+@app.on_event("startup")
+async def startup_event():
+    """Загрузка модели при старте"""
+    logger.info("🚀 BFS Face Swap API starting...")
+    success = load_pipeline()
+    if not success:
+        logger.warning("⚠️ Model failed to load, API will return fallback responses")
 
 @app.post("/swap")
 async def swap_face(
-    target: UploadFile = File(...),
-    source: UploadFile = File(...),
-    face_enhancer: str = Form("true"),
-    face_swapper_model: str = Form("inswapper_128"),
-    face_enhancer_model: str = Form("gfpgan_1.4")
+    target: UploadFile = File(...),  # BODY (тело) - ПЕРВОЕ!
+    source: UploadFile = File(...),  # FACE (лицо) - ВТОРОЕ!
+    num_steps: int = Form(20),  # Можно уменьшить до 10-15 для скорости
+    guidance_scale: float = Form(1.0)
 ):
-    """Замена лица – возвращает изображение"""
+    """
+    Замена лица с BFS Head V5
+    ВАЖНО: порядок файлов - сначала тело (target), потом лицо (source) [citation:2][citation:4]
+    """
     target_path = source_path = None
+    
     try:
+        # Сохраняем файлы
         target_path = save_upload_file(target)
         source_path = save_upload_file(source)
-        logger.info(f"🔄 Processing swap: {target.filename} <- {source.filename}")
-
-        result_arr = process_face_swap(
-            source_path, target_path,
-            swapper_model=face_swapper_model,
-            enhancer_model=face_enhancer_model if face_enhancer.lower() == "true" else None
-        )
-        result_bytes = write_numpy_image(result_arr, quality=98)
-        # Финальное улучшение
-        enhanced_bytes = enhance_image_quality(result_bytes)
-
-        logger.info(f"✅ Swap completed, size: {len(enhanced_bytes)} bytes")
-        return Response(content=enhanced_bytes, media_type="image/jpeg")
+        
+        logger.info(f"🔄 Processing BFS V5 swap: {target.filename} (body) <- {source.filename} (face)")
+        
+        # Загружаем и оптимизируем изображения
+        body_img = optimize_image(Image.open(target_path).convert("RGB"))
+        face_img = optimize_image(Image.open(source_path).convert("RGB"))
+        
+        if pipe is not None and lora_loaded:
+            # Промпт для Head V5 из официальной документации [citation:2][citation:4]
+            prompt = """head_swap: start with Picture 1 as the base image, keeping its lighting, environment, and background. remove the head from Picture 1 completely and replace it with the head from Picture 2, strictly preserving the hair, eye color, and nose structure of Picture 2. copy the eye direction, head rotation, and micro-expressions from Picture 1. high quality, sharp details, 4k"""
+            
+            # Инвертированный порядок: [body, face] [citation:2]
+            inputs = {
+                "image": [body_img, face_img],
+                "prompt": prompt,
+                "generator": torch.manual_seed(42),
+                "true_cfg_scale": 4.0,
+                "negative_prompt": "blurry, low quality, distorted face, bad anatomy, unnatural lighting",
+                "num_inference_steps": num_steps,
+                "guidance_scale": guidance_scale,
+            }
+            
+            logger.info("🔄 Running inference (this takes 1-3 minutes on CPU)...")
+            with torch.inference_mode():
+                output = pipe(**inputs)
+                result_img = output.images[0]
+        else:
+            # Fallback: если модель не загрузилась
+            logger.warning("Using fallback mode - returning body image")
+            result_img = body_img
+        
+        # Сохраняем результат с высоким качеством
+        result_bytes = io.BytesIO()
+        result_img.save(result_bytes, format="JPEG", quality=98, optimize=True)
+        
+        logger.info(f"✅ Swap completed: {len(result_bytes.getvalue())} bytes")
+        return Response(content=result_bytes.getvalue(), media_type="image/jpeg")
+        
     except Exception as e:
         logger.error(f"❌ Swap error: {e}")
         raise HTTPException(status_code=500, detail=str(e))
     finally:
+        # Очистка временных файлов
         for p in (target_path, source_path):
             if p and p.exists():
                 p.unlink()
 
-@app.post("/enhance")
-async def enhance_face(
-    image: UploadFile = File(...),
-    enhancer_model: str = Form("gfpgan_1.4")
+@app.post("/swap-with-prompt")
+async def swap_face_with_prompt(
+    target: UploadFile = File(...),  # BODY
+    source: UploadFile = File(...),  # FACE
+    custom_prompt: str = Form(...),  # Кастомный промпт
+    num_steps: int = Form(20)
 ):
-    """Улучшение лица – возвращает изображение"""
-    img_path = None
+    """Замена лица с кастомным промптом"""
+    target_path = source_path = None
+    
     try:
-        img_path = save_upload_file(image)
-        logger.info(f"🔧 Enhancing: {image.filename}")
-
-        result_arr = process_face_enhance(img_path, enhancer_model)
-        result_bytes = write_numpy_image(result_arr, quality=98)
-        enhanced_bytes = enhance_image_quality(result_bytes)
-
-        logger.info(f"✅ Enhancement completed, size: {len(enhanced_bytes)} bytes")
-        return Response(content=enhanced_bytes, media_type="image/jpeg")
+        target_path = save_upload_file(target)
+        source_path = save_upload_file(source)
+        
+        body_img = optimize_image(Image.open(target_path).convert("RGB"))
+        face_img = optimize_image(Image.open(source_path).convert("RGB"))
+        
+        if pipe is not None and lora_loaded:
+            # Добавляем требования качества к кастомному промпту
+            enhanced_prompt = f"{custom_prompt}. high quality, sharp details, 4k, photorealistic"
+            
+            inputs = {
+                "image": [body_img, face_img],
+                "prompt": enhanced_prompt,
+                "generator": torch.manual_seed(42),
+                "true_cfg_scale": 4.0,
+                "negative_prompt": "blurry, low quality, distorted",
+                "num_inference_steps": num_steps,
+                "guidance_scale": 1.0,
+            }
+            
+            with torch.inference_mode():
+                output = pipe(**inputs)
+                result_img = output.images[0]
+        else:
+            result_img = body_img
+        
+        result_bytes = io.BytesIO()
+        result_img.save(result_bytes, format="JPEG", quality=98)
+        
+        return Response(content=result_bytes.getvalue(), media_type="image/jpeg")
+        
     except Exception as e:
-        logger.error(f"❌ Enhance error: {e}")
+        logger.error(f"❌ Swap error: {e}")
         raise HTTPException(status_code=500, detail=str(e))
     finally:
-        if img_path and img_path.exists():
-            img_path.unlink()
+        for p in (target_path, source_path):
+            if p and p.exists():
+                p.unlink()
 
-@app.post("/analyse")
-async def analyse_face(image: UploadFile = File(...)):
-    """Анализ лиц – возвращает JSON"""
-    img_path = None
-    try:
-        img_path = save_upload_file(image)
-        analysis = process_face_analyse(img_path)
-        return JSONResponse(content={"status": "success", "analysis": analysis})
-    except Exception as e:
-        logger.error(f"❌ Analyse error: {e}")
-        raise HTTPException(status_code=500, detail=str(e))
-    finally:
-        if img_path and img_path.exists():
-            img_path.unlink()
+@app.get("/health")
+async def health():
+    """Проверка статуса"""
+    return {
+        "status": "ok",
+        "model_loaded": pipe is not None,
+        "lora_loaded": lora_loaded,
+        "base_model": BASE_MODEL,
+        "lora_file": LORA_FILE
+    }
 
 @app.get("/models")
-async def get_available_models():
-    """Список доступных моделей"""
+async def get_models():
+    """Информация о доступных версиях BFS [citation:2]"""
     return JSONResponse(content={
-        "face_swappers": [
-            {"name": "inswapper_128", "description": "Fast and accurate face swapper"},
-            {"name": "simswap", "description": "High quality face swapper"}
+        "base_model": BASE_MODEL,
+        "lora_repo": LORA_REPO,
+        "current_lora": LORA_FILE,
+        "available_versions": [
+            {
+                "version": "Face V1",
+                "file": "bfs_face_v1_qwen_image_edit_2509.safetensors",
+                "order": "Face then Body",
+                "description": "Swaps only face, preserves hair"
+            },
+            {
+                "version": "Head V1",
+                "file": "bfs_head_v1_qwen_image_edit_2509.safetensors",
+                "order": "Face then Body",
+                "description": "Full head swap"
+            },
+            {
+                "version": "Head V3 (Recommended for 2509)",
+                "file": "bfs_head_v3_qwen_image_edit_2509.safetensors",
+                "order": "Body then Face",
+                "description": "Most stable for 2509"
+            },
+            {
+                "version": "Head V5 (Recommended for 2511)",
+                "file": "bfs_head_v5_2511_merged_version_rank_16_fp16.safetensors",
+                "order": "Body then Face",
+                "description": "Latest, best expression transfer"
+            }
         ],
-        "face_enhancers": [
-            {"name": "gfpgan_1.4", "description": "Face restoration model"},
-            {"name": "codeformer", "description": "Face enhancement model"}
-        ],
-        "face_detectors": [
-            {"name": "retinaface", "description": "Accurate face detector"},
-            {"name": "mtcnn", "description": "Multi-task face detector"}
-        ]
+        "prompts": {
+            "head_v5": "head_swap: start with Picture 1 as the base image, keeping its lighting, environment, and background. remove the head from Picture 1 completely and replace it with the head from Picture 2, strictly preserving the hair, eye color, and nose structure of Picture 2. copy the eye direction, head rotation, and micro-expressions from Picture 1. high quality, sharp details, 4k"
+        }
     })
 
-@app.on_event("startup")
-async def startup_event():
-    logger.info("🚀 FaceFusion API starting...")
-    if FACEFUSION_AVAILABLE:
-        logger.info("✅ FaceFusion loaded successfully")
-    else:
-        logger.warning(f"⚠️ FaceFusion not available, using mock mode. Error: {facefusion_import_error}")
-
 @app.on_event("shutdown")
 async def shutdown_event():
+    """Очистка при остановке"""
     logger.info("🛑 Shutting down...")
     shutil.rmtree(TEMP_DIR, ignore_errors=True)