Create controlnet_facefix.py

controlnet_facefix.py

@@ -0,0 +1,68 @@
+import torch
+from diffusers import StableDiffusionControlNetInpaintPipeline, ControlNetModel
+from controlnet_aux import OpenposeDetector, ZoeDetector
+from PIL import Image
+
+# ───── Globale Modelle (einmal laden, bleibt im VRAM) ─────
+print("Lade OpenPose_faceonly + Depth für perfekte Gesichter/Hände...")
+
+# OpenPose Face-Only Preprocessor + ControlNet
+openpose_face = OpenposeDetector.from_pretrained(
+    "lllyasviel/ControlNet", model_name="openpose_face"
+)
+
+# Depth Preprocessor + ControlNet
+depth_processor = ZoeDetector.from_pretrained("lllyasviel/ControlNet")
+
+controlnet_face = ControlNetModel.from_pretrained(
+    "lllyasviel/control_v11p_sd15_openpose",
+    subfolder="faceonly",
+    torch_dtype=torch.float16
+).to("cuda")
+
+controlnet_depth = ControlNetModel.from_pretrained(
+    "lllyasviel/control_v11f1e_sd15_depth",
+    torch_dtype=torch.float16
+).to("cuda")
+
+# Pipeline-Cache (wird erst beim ersten Aufruf erstellt)
+_facefix_pipe = None
+
+def _get_facefix_pipeline(model_id: str):
+    global _facefix_pipe
+    if _facefix_pipe is None:
+        print(f"Lade Face-Fix-Pipeline mit Modell: {model_id}")
+        _facefix_pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(
+            model_id,
+            controlnet=[controlnet_face, controlnet_depth],
+            torch_dtype=torch.float16,
+            safety_checker=None,
+            requires_safety_checker=False,
+        ).to("cuda")
+        _facefix_pipe.enable_xformers_memory_efficient_attention()
+        _facefix_pipe.enable_model_cpu_offload()  # spart ~2 GB bei 16 GB Karten
+    return _facefix_pipe
+
+def apply_facefix(image: Image.Image, prompt: str, negative_prompt: str, seed: int, base_model_path: str):
+    """
+    Automatischer 20-Sekunden-Fix für perfekte Gesichter & Hände
+    """
+    pipe = _get_facefix_pipeline(base_model_path)
+
+    # Control-Images aus dem generierten Bild erzeugen
+    pose_img = openpose_face(image)
+    depth_img = depth_processor(image)
+
+    fixed_image = pipe(
+        prompt=prompt,
+        negative_prompt=negative_prompt,
+        image=image,
+        control_image=[pose_img, depth_img],
+        controlnet_conditioning_scale=[0.85, 0.60],   # Face stark, Depth mittel
+        strength=0.42,
+        num_inference_steps=20,
+        guidance_scale=7.0,
+        generator=torch.Generator("cuda").manual_seed(seed),
+    ).images[0]
+
+    return fixed_image