Update app.py

app.py

@@ -1,145 +1,131 @@
 import os
-import subprocess
+import torch
+import uvicorn
 import tempfile
 import uuid
-import json
 import time
-import requests
-from fastapi import FastAPI, File, UploadFile, HTTPException
-from fastapi.responses import Response, JSONResponse
-import uvicorn
 import logging
+from fastapi import FastAPI, File, UploadFile, HTTPException
+from fastapi.responses import Response
+from PIL import Image
+from diffusers import QwenImageEditPlusPipeline
 
+# Настройка логирования
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-app = FastAPI()
+app = FastAPI(title="Head Swap API (Qwen + BFS LoRA)")
 
-# Запускаем ComfyUI в фоновом режиме
-comfy_process = None
+# Глобальные переменные для пайплайна
+pipe = None
+device = "cpu"  # Работаем на CPU
 
-@app.on_event("startup")
-async def startup_event():
-    global comfy_process
-    logger.info("Starting ComfyUI...")
-    comfy_process = subprocess.Popen(
-        ["python", "/comfyui/main.py", "--listen", "0.0.0.0", "--port", "8188"],
-        stdout=subprocess.PIPE,
-        stderr=subprocess.PIPE
-    )
-    # Даём время на запуск
-    time.sleep(10)
-    logger.info("ComfyUI started")
+# Константы для генерации
+NUM_INFERENCE_STEPS = 40
+TRUE_GUIDANCE_SCALE = 4.0
+NEGATIVE_PROMPT = " "  # Пустой негативный промпт (как в оригинале)
 
-@app.on_event("shutdown")
-async def shutdown_event():
-    if comfy_process:
-        comfy_process.terminate()
+# Фиксированный промпт для замены головы (следуя рекомендациям BFS V3)
+# Picture 1 – целевое изображение (тело), Picture 2 – изображение лица (источник)
+HEAD_SWAP_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, "
+    "ensuring a seamless and natural blend. maintain the facial identity, expression, and features from Picture 2."
+)
 
-def load_workflow():
-    """Загружает workflow из файла или создаёт базовый"""
-    # Здесь можно загрузить ваш workflow.json
-    # Пока используем простой workflow
-    return {
-        "3": {
-            "class_type": "LoadImage",
-            "inputs": {
-                "image": "source.jpg"
-            }
-        },
-        "4": {
-            "class_type": "LoadImage",
-            "inputs": {
-                "image": "target.jpg"
-            }
-        },
-        "5": {
-            "class_type": "ReActorFaceSwap",
-            "inputs": {
-                "source_image": ["3", 0],
-                "target_image": ["4", 0],
-                "face_restorer": "gfpgan",
-                "face_restorer_weight": 0.8,
-                "swap_model": "inswapper_128.onnx",
-                "detect_model": "yolov8n-face.pt",
-                "save_original": False,
-                "output_image": ["6", 0]
-            }
-        },
-        "6": {
-            "class_type": "SaveImage",
-            "inputs": {
-                "filename_prefix": "output",
-                "images": ["5", 0]
-            }
-        }
-    }
+@app.on_event("startup")
+async def load_model():
+    global pipe
+    logger.info("Loading Qwen-Image-Edit-2511 model...")
+    # Загружаем пайплайн с float32 для CPU
+    pipe = QwenImageEditPlusPipeline.from_pretrained(
+        "Qwen/Qwen-Image-Edit-2511",
+        torch_dtype=torch.float32,
+        safety_checker=None  # отключаем safety checker для скорости
+    )
+    pipe = pipe.to(device)
+    
+    # Загружаем LoRA BFS Head V3
+    lora_path = "/app/bfs_head_v3_qwen_image_edit_2509.safetensors"
+    if os.path.exists(lora_path):
+        logger.info("Loading BFS LoRA weights...")
+        pipe.load_lora_weights(lora_path, adapter_name="bfs")
+        pipe.set_adapter("bfs")
+    else:
+        logger.warning("LoRA file not found, proceeding without it.")
+    
+    logger.info("Model ready.")
 
 @app.post("/swap")
-async def swap_face(
-    source: UploadFile = File(...),
-    target: UploadFile = File(...)
+async def swap_head(
+    target: UploadFile = File(..., description="Target image (body)"),
+    source: UploadFile = File(..., description="Source image (face)")
 ):
-    temp_dir = tempfile.mkdtemp()
+    """
+    Заменяет голову на целевом изображении (target) лицом из source.
+    Порядок важен: target – тело, source – лицо (BFS V3 инвертированный порядок).
+    """
+    temp_dir = None
     try:
-        # Сохраняем входные файлы
-        source_path = os.path.join(temp_dir, "source.jpg")
-        target_path = os.path.join(temp_dir, "target.jpg")
+        # Сохраняем загруженные файлы во временную директорию
+        temp_dir = tempfile.mkdtemp()
+        target_path = os.path.join(temp_dir, f"target_{uuid.uuid4().hex}.jpg")
+        source_path = os.path.join(temp_dir, f"source_{uuid.uuid4().hex}.jpg")
         
-        with open(source_path, "wb") as f:
-            f.write(await source.read())
         with open(target_path, "wb") as f:
             f.write(await target.read())
-
-        # Копируем файлы в папку ComfyUI
-        comfy_input_dir = "/comfyui/input"
-        os.makedirs(comfy_input_dir, exist_ok=True)
+        with open(source_path, "wb") as f:
+            f.write(await source.read())
         
-        shutil.copy(source_path, os.path.join(comfy_input_dir, "source.jpg"))
-        shutil.copy(target_path, os.path.join(comfy_input_dir, "target.jpg"))
-
-        # Загружаем workflow
-        workflow = load_workflow()
+        # Открываем изображения как PIL
+        target_img = Image.open(target_path).convert("RGB")
+        source_img = Image.open(source_path).convert("RGB")
+        
+        # Пайплайн ожидает список изображений: [target, source] (в таком порядке)
+        input_images = [target_img, source_img]
+        
+        logger.info("Starting generation...")
+        start_time = time.time()
+        
+        # Генерация
+        result_images = pipe(
+            image=input_images,
+            prompt=HEAD_SWAP_PROMPT,
+            negative_prompt=NEGATIVE_PROMPT,
+            num_inference_steps=NUM_INFERENCE_STEPS,
+            true_cfg_scale=TRUE_GUIDANCE_SCALE,
+            generator=torch.Generator(device=device).manual_seed(42),  # фиксированный seed для воспроизводимости
+        ).images
+        
+        elapsed = time.time() - start_time
+        logger.info(f"Generation took {elapsed:.2f} seconds")
         
-        # Отправляем задачу в ComfyUI
-        response = requests.post(
-            "http://127.0.0.1:8188/prompt",
-            json={"prompt": workflow}
-        )
+        if not result_images:
+            raise HTTPException(500, "No image generated")
         
-        if response.status_code != 200:
-            raise HTTPException(500, f"ComfyUI error: {response.text}")
+        # Сохраняем результат во временный файл и возвращаем
+        result_img = result_images[0]
+        output_path = os.path.join(temp_dir, f"output_{uuid.uuid4().hex}.jpg")
+        result_img.save(output_path, format="JPEG")
         
-        prompt_id = response.json()["prompt_id"]
+        with open(output_path, "rb") as f:
+            image_data = f.read()
         
-        # Ждём завершения
-        while True:
-            status = requests.get(f"http://127.0.0.1:8188/history/{prompt_id}")
-            if status.status_code == 200 and status.json():
-                history = status.json()
-                if prompt_id in history:
-                    output_images = history[prompt_id]["outputs"]
-                    # Находим выходное изображение
-                    for node_id, node_output in output_images.items():
-                        if "images" in node_output:
-                            image_info = node_output["images"][0]
-                            image_path = os.path.join("/comfyui/output", image_info["filename"])
-                            if os.path.exists(image_path):
-                                with open(image_path, "rb") as f:
-                                    image_data = f.read()
-                                return Response(content=image_data, media_type="image/jpeg")
-            time.sleep(1)
-            
+        return Response(content=image_data, media_type="image/jpeg")
+    
     except Exception as e:
-        logger.exception("Error")
+        logger.exception("Error during head swap")
         raise HTTPException(500, str(e))
+    
     finally:
-        import shutil
-        shutil.rmtree(temp_dir, ignore_errors=True)
+        # Очистка временных файлов
+        if temp_dir and os.path.exists(temp_dir):
+            import shutil
+            shutil.rmtree(temp_dir, ignore_errors=True)
 
 @app.get("/health")
 async def health():
-    return {"status": "ok"}
+    return {"status": "ok", "device": device, "model": "Qwen-Image-Edit-2511 with BFS LoRA"}
 
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)