import torch
from diffusers import StableDiffusionControlNetInpaintPipeline, ControlNetModel
from controlnet_aux import ZoeDetector
from PIL import Image
import time

print("Lade Face-Fix System (Depth-basiert)...")

# Vereinfachte Version NUR mit Depth
try:
    depth_processor = ZoeDetector.from_pretrained("lllyasviel/ControlNet")
    print("✅ Depth Processor geladen")
    
    controlnet_depth = ControlNetModel.from_pretrained(
        "lllyasviel/control_v11f1e_sd15_depth",
        torch_dtype=torch.float16
    ).to("cuda")
    print("✅ ControlNet Depth geladen")
    
    _facefix_pipe = None
    
    def _get_facefix_pipeline(model_id: str):
        global _facefix_pipe
        if _facefix_pipe is None:
            print(f"Lade Face-Fix-Pipeline für Modell: {model_id}")
            _facefix_pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(
                model_id,
                controlnet=[controlnet_depth],
                torch_dtype=torch.float16,
                safety_checker=None,
            ).to("cuda")
        return _facefix_pipe
    
    def apply_facefix(image: Image.Image, prompt: str, negative_prompt: str, seed: int, model_id: str):
        try:
            pipe = _get_facefix_pipeline(model_id)
            depth_img = depth_processor(image)
            
            face_prompt = prompt + ", perfect face, detailed skin"
            face_negative = negative_prompt + ", deformed face"
            
            fixed = pipe(
                prompt=face_prompt,
                negative_prompt=face_negative,
                image=image,
                mask_image=None,
                control_image=[depth_img],
                controlnet_conditioning_scale=[0.7],
                strength=0.35,
                num_inference_steps=20,
                guidance_scale=7.0,
                generator=torch.Generator("cuda").manual_seed(seed),
            ).images[0]
            
            return fixed
            
        except Exception as e:
            print(f"Face-Fix Error: {e}")
            return image
            
except Exception as e:
    print(f"❌ Face-Fix Setup fehlgeschlagen: {e}")
    
    # Fallback: Dummy-Funktion die nichts macht
    def apply_facefix(image: Image.Image, prompt: str, negative_prompt: str, seed: int, model_id: str):
        print("⚠️ Face-Fix nicht verfügbar, gebe Original zurück")
        return image