app.py

"""
Pro AI Headshot Generator - Public Access Version
Transforms any selfie into professional headshots using AI.
"""
import cv2
import torch
import random
import numpy as np
import os
import time
import glob
from pathlib import Path
from typing import Tuple, Optional
from datetime import datetime, timedelta

try:
    import spaces
    SPACES_AVAILABLE = True
except ImportError:
    SPACES_AVAILABLE = False
    # Create a dummy decorator for local development
    class spaces:
        @staticmethod
        def GPU(func):
            return func

import PIL
from PIL import Image

import diffusers
from diffusers.utils import load_image
from diffusers.models import ControlNetModel
from diffusers.pipelines.controlnet.multicontrolnet import MultiControlNetModel
from huggingface_hub import hf_hub_download
from insightface.app import FaceAnalysis

from style_template import styles
from pipeline_stable_diffusion_xl_instantid_full import (
    StableDiffusionXLInstantIDPipeline,
    draw_kps,
)
import gradio as gr
import json
import sys
import gc
import os

# Ensure the current directory and parent directory are in Python path (for Hugging Face Spaces)
current_dir = os.path.dirname(os.path.abspath(__file__))
parent_dir = os.path.dirname(current_dir)
for dir_path in [current_dir, parent_dir]:
    if dir_path not in sys.path:
        sys.path.insert(0, dir_path)

from depth_anything.dpt import DepthAnything
from depth_anything.util.transform import Resize, NormalizeImage, PrepareForNet
import torch.nn.functional as F
from torchvision.transforms import Compose

# Try to import config, with fallback values if it fails
try:
    from config import (
        TEMP_DIR,
        MAX_FILE_AGE_HOURS,
        MAX_IMAGE_SIZE_MB,
        ALLOWED_IMAGE_FORMATS,
        MIN_IMAGE_DIMENSION,
        MAX_IMAGE_DIMENSION,
        MAX_PROMPT_LENGTH,
        MAX_NEGATIVE_PROMPT_LENGTH,
        DEFAULT_NUM_STEPS,
        DEFAULT_GUIDANCE_SCALE,
        DEFAULT_SEED,
    )
except ImportError:
    # Fallback configuration if config.py is not found
    from pathlib import Path
    TEMP_DIR = Path(os.getenv("TEMP_DIR", "temp_downloads"))
    TEMP_DIR.mkdir(exist_ok=True, parents=True)
    MAX_FILE_AGE_HOURS = int(os.getenv("MAX_FILE_AGE_HOURS", "24"))
    MAX_IMAGE_SIZE_MB = int(os.getenv("MAX_IMAGE_SIZE_MB", "10"))
    ALLOWED_IMAGE_FORMATS = {".jpg", ".jpeg", ".png", ".webp"}
    MIN_IMAGE_DIMENSION = 128
    MAX_IMAGE_DIMENSION = 4096
    DEFAULT_NUM_STEPS = int(os.getenv("DEFAULT_NUM_STEPS", "30"))
    DEFAULT_GUIDANCE_SCALE = float(os.getenv("DEFAULT_GUIDANCE_SCALE", "5.0"))
    DEFAULT_SEED = int(os.getenv("DEFAULT_SEED", "42"))
    MAX_PROMPT_LENGTH = int(os.getenv("MAX_PROMPT_LENGTH", "500"))
    MAX_NEGATIVE_PROMPT_LENGTH = int(os.getenv("MAX_NEGATIVE_PROMPT_LENGTH", "500"))
    print("⚠ Warning: config.py not found, using default configuration values")


# ============ GLOBAL CONFIG ============
MAX_SEED = np.iinfo(np.int32).max

# Device detection - fixed logic
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float16 if torch.cuda.is_available() else torch.float32
STYLE_NAMES = list(styles.keys())


# ============ ZEROGPU HELPER FUNCTIONS ============
def serialize_for_zerogpu(obj):
    """
    Serialize objects for ZeroGPU return values.
    Converts numpy arrays, pandas DataFrames, and other non-serializable objects
    to JSON-serializable formats.
    """
    try:
        import pandas as pd
        HAS_PANDAS = True
    except ImportError:
        HAS_PANDAS = False
    
    if isinstance(obj, np.ndarray):
        return obj.tolist()
    elif isinstance(obj, (np.integer, np.floating)):
        return float(obj)
    elif isinstance(obj, (np.bool_)):
        return bool(obj)
    elif HAS_PANDAS and isinstance(obj, pd.DataFrame):
        # Convert DataFrame to dict with limited rows to avoid large returns
        max_rows = 1000  # Limit to prevent large returns
        df_limited = obj.head(max_rows) if len(obj) > max_rows else obj
        return {
            'data': df_limited.to_dict('records'),
            'columns': df_limited.columns.tolist(),
            'total_rows': len(obj),
            'limited': len(obj) > max_rows
        }
    elif HAS_PANDAS and isinstance(obj, pd.Series):
        return obj.tolist()
    elif isinstance(obj, dict):
        return {k: serialize_for_zerogpu(v) for k, v in obj.items()}
    elif isinstance(obj, (list, tuple)):
        return [serialize_for_zerogpu(item) for item in obj]
    elif isinstance(obj, (str, int, float, bool, type(None))):
        return obj
    else:
        # Try to convert to string as fallback
        try:
            return str(obj)
        except:
            return None


def safe_zerogpu_return(*args, **kwargs):
    """
    Wrapper to safely return values from ZeroGPU functions.
    Ensures all return values are serializable and not too large.
    """
    results = []
    if args:
        if len(args) == 1:
            results = serialize_for_zerogpu(args[0])
        else:
            results = [serialize_for_zerogpu(arg) for arg in args]
    
    # Clean up GPU memory before returning
    if device == "cuda":
        try:
            torch.cuda.empty_cache()
        except:
            pass
    gc.collect()
    
    return results


# ============ FILE MANAGEMENT ============
def cleanup_old_files():
    """Remove files older than MAX_FILE_AGE_HOURS from temp directory."""
    if not TEMP_DIR.exists():
        return
    
    cutoff_time = time.time() - (MAX_FILE_AGE_HOURS * 3600)
    deleted_count = 0
    
    for file_path in TEMP_DIR.glob("*"):
        if file_path.is_file():
            try:
                if file_path.stat().st_mtime < cutoff_time:
                    file_path.unlink()
                    deleted_count += 1
            except Exception as e:
                print(f"Warning: Could not delete {file_path}: {e}")
    
    if deleted_count > 0:
        print(f"Cleaned up {deleted_count} old files from temp directory")


def save_as_png(image: Image.Image, filename: str = "professional_headshot") -> str:
    """Save image as PNG in temp directory with cleanup."""
    # Cleanup old files before saving new ones
    cleanup_old_files()
    
    TEMP_DIR.mkdir(exist_ok=True, parents=True)
    timestamp = int(time.time())
    filepath = TEMP_DIR / f"{filename}_{timestamp}.png"

    # Ensure image is in RGB mode
    if image.mode in ("RGBA", "LA"):
        background = Image.new("RGB", image.size, (255, 255, 255))
        background.paste(image, mask=image.split()[-1])
        image = background
    elif image.mode != "RGB":
        image = image.convert("RGB")

    image.save(filepath, "PNG", optimize=True)
    return str(filepath)


# ============ INPUT VALIDATION ============
def validate_image_file(file_path: str) -> Tuple[bool, Optional[str]]:
    """Validate uploaded image file."""
    if not file_path:
        return False, "Please upload an image file."
    
    # Check file exists
    if not os.path.exists(file_path):
        return False, "Uploaded file not found. Please try again."
    
    # Check file extension
    ext = Path(file_path).suffix.lower()
    if ext not in ALLOWED_IMAGE_FORMATS:
        return False, f"Unsupported file format. Allowed formats: {', '.join(ALLOWED_IMAGE_FORMATS)}"
    
    # Check file size
    file_size_mb = os.path.getsize(file_path) / (1024 * 1024)
    if file_size_mb > MAX_IMAGE_SIZE_MB:
        return False, f"File too large. Maximum size: {MAX_IMAGE_SIZE_MB}MB"
    
    # Try to open and validate image
    try:
        with Image.open(file_path) as img:
            width, height = img.size
            if width < MIN_IMAGE_DIMENSION or height < MIN_IMAGE_DIMENSION:
                return False, f"Image too small. Minimum size: {MIN_IMAGE_DIMENSION}x{MIN_IMAGE_DIMENSION}px"
            if width > MAX_IMAGE_DIMENSION or height > MAX_IMAGE_DIMENSION:
                return False, f"Image too large. Maximum size: {MAX_IMAGE_DIMENSION}x{MAX_IMAGE_DIMENSION}px"
    except Exception as e:
        return False, f"Invalid image file: {str(e)}"
    
    return True, None


def validate_prompt(prompt: str) -> Tuple[bool, Optional[str]]:
    """Validate prompt input."""
    if not prompt or not prompt.strip():
        return True, None  # Empty prompt is allowed (will use default)
    
    if len(prompt) > MAX_PROMPT_LENGTH:
        return False, f"Prompt too long. Maximum length: {MAX_PROMPT_LENGTH} characters"
    
    return True, None


def validate_negative_prompt(negative_prompt: str) -> Tuple[bool, Optional[str]]:
    """Validate negative prompt input."""
    if not negative_prompt:
        return True, None
    
    if len(negative_prompt) > MAX_NEGATIVE_PROMPT_LENGTH:
        return False, f"Negative prompt too long. Maximum length: {MAX_NEGATIVE_PROMPT_LENGTH} characters"
    
    return True, None


# ============ MODEL LOADING ============
print("Loading AI models... This may take a few minutes on first run.")
print("=" * 60)

def check_file_exists(local_path):
    """Check if a file exists locally."""
    return Path(local_path).exists()

try:
    # InstantID checkpoints - only download if missing
    print("Checking InstantID checkpoints...")
    checkpoint_dir = Path("./checkpoints/ControlNetModel")
    config_path = checkpoint_dir / "config.json"
    model_path = checkpoint_dir / "diffusion_pytorch_model.safetensors"
    adapter_path = Path("./checkpoints/ip-adapter.bin")
    
    if not config_path.exists():
        print("  Downloading config.json...")
        hf_hub_download(
            repo_id="InstantX/InstantID",
            filename="ControlNetModel/config.json",
            local_dir="./checkpoints",
        )
        print("  ✓ config.json downloaded")
    else:
        print("  ✓ config.json already exists")
    
    if not model_path.exists():
        print("  Downloading diffusion_pytorch_model.safetensors (this may take a while)...")
        hf_hub_download(
            repo_id="InstantX/InstantID",
            filename="ControlNetModel/diffusion_pytorch_model.safetensors",
            local_dir="./checkpoints",
        )
        print("  ✓ diffusion_pytorch_model.safetensors downloaded")
    else:
        print("  ✓ diffusion_pytorch_model.safetensors already exists")
    
    if not adapter_path.exists():
        print("  Downloading ip-adapter.bin...")
        hf_hub_download(
            repo_id="InstantX/InstantID", 
            filename="ip-adapter.bin", 
            local_dir="./checkpoints"
        )
        print("  ✓ ip-adapter.bin downloaded")
    else:
        print("  ✓ ip-adapter.bin already exists")
    
    print("✓ InstantID checkpoints ready")
except Exception as e:
    print(f"⚠ Warning: Error downloading InstantID checkpoints: {e}")
    print("  The app will attempt to use cached models if available.")
    import traceback
    traceback.print_exc()

try:
    # Face encoder
    print("\nLoading face recognition model...")
    models_dir = Path("./models/antelopev2")
    if not models_dir.exists():
        print("  Warning: Face models directory not found. Models will be downloaded automatically.")
    
    app = FaceAnalysis(
        name="antelopev2",
        root="./",
        providers=["CPUExecutionProvider"],
    )
    app.prepare(ctx_id=0, det_size=(640, 640))
    print("✓ Face recognition model loaded")
except Exception as e:
    print(f"✗ Error loading face recognition model: {e}")
    import traceback
    traceback.print_exc()
    raise RuntimeError("Failed to load face recognition model. Please ensure models are downloaded correctly.")

try:
    # DepthAnything
    print("\nLoading depth estimation model...")
    print("  This may take a few minutes on first run (downloading ~1.5GB)...")
    depth_anything = DepthAnything.from_pretrained(
        "LiheYoung/depth_anything_vitl14"
    ).to(device).eval()
    print("✓ Depth estimation model loaded")
except Exception as e:
    print(f"✗ Error loading depth estimation model: {e}")
    import traceback
    traceback.print_exc()
    raise RuntimeError("Failed to load depth estimation model. Please check your internet connection and try again.")

transform = Compose(
    [
        Resize(
            width=518,
            height=518,
            resize_target=False,
            keep_aspect_ratio=True,
            ensure_multiple_of=14,
            resize_method="lower_bound",
            image_interpolation_method=cv2.INTER_CUBIC,
        ),
        NormalizeImage(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        PrepareForNet(),
    ]
)

face_adapter = "./checkpoints/ip-adapter.bin"
controlnet_path = "./checkpoints/ControlNetModel"

try:
    print("\nLoading ControlNet models...")
    print("  Loading Identity ControlNet...")
    controlnet_identitynet = ControlNetModel.from_pretrained(
        controlnet_path, torch_dtype=dtype
    )
    print("  ✓ Identity ControlNet loaded")
    
    controlnet_canny_model = "diffusers/controlnet-canny-sdxl-1.0"
    controlnet_depth_model = "diffusers/controlnet-depth-sdxl-1.0-small"
    
    print("  Loading Canny ControlNet (this may take a while)...")
    controlnet_canny = ControlNetModel.from_pretrained(
        controlnet_canny_model, torch_dtype=dtype
    ).to(device)
    print("  ✓ Canny ControlNet loaded")
    
    print("  Loading Depth ControlNet (this may take a while)...")
    controlnet_depth = ControlNetModel.from_pretrained(
        controlnet_depth_model, torch_dtype=dtype
    ).to(device)
    print("  ✓ Depth ControlNet loaded")
    print("✓ All ControlNet models loaded")
except Exception as e:
    print(f"✗ Error loading ControlNet models: {e}")
    import traceback
    traceback.print_exc()
    raise RuntimeError("Failed to load ControlNet models. Please check your internet connection and disk space.")


def get_depth_map(image):
    """Generate depth map from image."""
    try:
        print("    Processing depth estimation...")
        image = np.array(image) / 255.0
        h, w = image.shape[:2]
        print(f"    Input size: {w}x{h}")
        
        print("    Applying transforms...")
        image = transform({"image": image})["image"]
        image = torch.from_numpy(image).unsqueeze(0).to(device)
        
        print("    Running depth model (this may take a while on CPU)...")
        with torch.no_grad():
            depth = depth_anything(image)
        
        print("    Post-processing depth map...")
        depth = F.interpolate(
            depth[None], (h, w), mode="bilinear", align_corners=False
        )[0, 0]
        depth = (depth - depth.min()) / (depth.max() - depth.min()) * 255.0
        depth = depth.cpu().numpy().astype(np.uint8)
        depth_image = Image.fromarray(depth)
        print("    ✓ Depth map generated")
        return depth_image
    except Exception as e:
        print(f"    ✗ Error generating depth map: {e}")
        import traceback
        traceback.print_exc()
        # Return a blank depth map as fallback
        return Image.new("L", (512, 512), 128)


def get_canny_image(image, t1=100, t2=200):
    """Generate canny edge map from image."""
    try:
        image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
        edges = cv2.Canny(image, t1, t2)
        return Image.fromarray(edges, "L")
    except Exception as e:
        print(f"Warning: Error generating canny edges: {e}")
        # Return a blank canny map as fallback
        return Image.new("L", (512, 512), 0)


controlnet_map = {
    "canny": controlnet_canny,
    "depth": controlnet_depth,
}
controlnet_map_fn = {
    "canny": get_canny_image,
    "depth": get_depth_map,
}

pretrained_model_name_or_path = "wangqixun/YamerMIX_v8"

try:
    print("\nLoading Stable Diffusion XL pipeline...")
    print("  This is the largest model (~6-7GB) and may take 10-20 minutes on first run...")
    print("  Please be patient and ensure you have a stable internet connection.")
    pipe = StableDiffusionXLInstantIDPipeline.from_pretrained(
        pretrained_model_name_or_path,
        controlnet=[controlnet_identitynet],
        torch_dtype=dtype,
        safety_checker=None,
        feature_extractor=None,
    ).to(device)
    print("✓ Stable Diffusion XL pipeline loaded")
    
    # Standard scheduler
    pipe.scheduler = diffusers.EulerDiscreteScheduler.from_config(
        pipe.scheduler.config
    )
    
    print("\nLoading IP-Adapter...")
    if device == "cuda":
        pipe.cuda()
        pipe.load_ip_adapter_instantid(face_adapter)
        # Ensure all components are on CUDA
        pipe.image_proj_model.to("cuda")
        pipe.unet.to("cuda")
        # Verify models have data (not meta tensors)
        test_param = next(pipe.unet.parameters())
        if hasattr(test_param, 'is_meta') and test_param.is_meta:
            print("  ⚠ Warning: UNet appears to have meta tensors")
        else:
            print(f"  ✓ UNet verified on device: {test_param.device}")
    else:
        pipe.load_ip_adapter_instantid(face_adapter)
        # Verify CPU models have data
        test_param = next(pipe.unet.parameters())
        if hasattr(test_param, 'is_meta') and test_param.is_meta:
            print("  ⚠ Warning: UNet appears to have meta tensors")
        else:
            print(f"  ✓ UNet verified on device: {test_param.device}")
    print("✓ IP-Adapter loaded")
    
    print("\n" + "=" * 60)
    print("✅ ALL MODELS LOADED SUCCESSFULLY!")
    print("=" * 60)
    print("The application is now ready to use.")
except Exception as e:
    print(f"\n✗ Error loading Stable Diffusion pipeline: {e}")
    import traceback
    traceback.print_exc()
    raise RuntimeError(f"Failed to load Stable Diffusion pipeline: {e}. Please check your internet connection and disk space.")


# ============ UTILS ============
def toggle_lcm_ui(value: bool) -> Tuple[dict, dict]:
    """Toggle UI for LCM mode."""
    if value:
        return (
            gr.update(minimum=0, maximum=100, step=1, value=5),
            gr.update(minimum=0.1, maximum=20.0, step=0.1, value=1.5),
        )
    else:
        return (
            gr.update(minimum=5, maximum=100, step=1, value=DEFAULT_NUM_STEPS),
            gr.update(minimum=0.1, maximum=20.0, step=0.1, value=DEFAULT_GUIDANCE_SCALE),
        )


def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
    """Randomize seed if requested."""
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    return seed


def convert_from_cv2_to_image(img: np.ndarray) -> Image.Image:
    """Convert OpenCV image to PIL Image."""
    return Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))


def convert_from_image_to_cv2(img: Image.Image) -> np.ndarray:
    """Convert PIL Image to OpenCV format."""
    return cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)


def resize_img(
    input_image,
    max_side=1280,
    min_side=1024,
    size=None,
    pad_to_max_side=False,
    mode=PIL.Image.BILINEAR,
    base_pixel_number=64,
):
    """Resize image maintaining aspect ratio."""
    w, h = input_image.size
    if size is not None:
        w_resize_new, h_resize_new = size
    else:
        ratio = min_side / min(h, w)
        w, h = round(ratio * w), round(ratio * h)
        ratio = max_side / max(h, w)
        input_image = input_image.resize([round(ratio * w), round(ratio * h)], mode)
        w_resize_new = (round(ratio * w) // base_pixel_number) * base_pixel_number
        h_resize_new = (round(ratio * h) // base_pixel_number) * base_pixel_number
    input_image = input_image.resize([w_resize_new, h_resize_new], mode)

    if pad_to_max_side:
        res = np.ones([max_side, max_side, 3], dtype=np.uint8) * 255
        offset_x = (max_side - w_resize_new) // 2
        offset_y = (max_side - h_resize_new) // 2
        res[offset_y : offset_y + h_resize_new, offset_x : offset_x + w_resize_new] = (
            np.array(input_image)
        )
        input_image = Image.fromarray(res)
    return input_image


def apply_style(style_name: str, positive: str, negative: str = "") -> Tuple[str, str]:
    """Apply style template to prompts."""
    if style_name == "No Style":
        return positive, negative
    p, n = styles.get(style_name, ("{prompt}", ""))
    return p.replace("{prompt}", positive), n + " " + negative


# ============ GENERATION FUNCTION ============
@spaces.GPU  # ZeroGPU will allocate GPU for this function (Hugging Face Spaces only)
def generate_image(
    face_image_path: str,
    prompt: str,
    negative_prompt: str,
    style_name: str,
    num_steps: int,
    identitynet_strength_ratio: float,
    adapter_strength_ratio: float,
    canny_strength: float,
    depth_strength: float,
    controlnet_selection: list,
    guidance_scale: float,
    seed: int,
    scheduler: str,
    enable_LCM: bool,
    enhance_face_region: bool,
    progress=gr.Progress(track_tqdm=True),
):
    """Generate professional headshot from face image."""
    try:
        print("\n" + "=" * 60)
        print("Starting image generation...")
        print("=" * 60)
        # Validate inputs
        print("Step 1/8: Validating inputs...")
        is_valid, error_msg = validate_image_file(face_image_path)
        if not is_valid:
            raise gr.Error(error_msg)
        
        is_valid, error_msg = validate_prompt(prompt)
        if not is_valid:
            raise gr.Error(error_msg)
        
        is_valid, error_msg = validate_negative_prompt(negative_prompt)
        if not is_valid:
            raise gr.Error(error_msg)
        print("  ✓ Inputs validated")

        # Randomize seed if needed
        if seed < 0:
            seed = random.randint(0, MAX_SEED)

        # Load and process face image
        print("Step 2/8: Loading and processing face image...")
        face_image = load_image(face_image_path)
        face_image = resize_img(face_image, max_side=1024)
        face_image_cv2 = convert_from_image_to_cv2(face_image)
        height, width, _ = face_image_cv2.shape
        print(f"  ✓ Image loaded: {width}x{height}")

        # Detect face
        print("Step 3/8: Detecting face...")
        face_info_list = app.get(face_image_cv2)
        if len(face_info_list) == 0:
            raise gr.Error(
                "Unable to detect a face in the image. Please upload a different photo with a clear face."
            )
        print(f"  ✓ Face detected")

        # Use largest detected face
        print("Step 4/8: Processing face features...")
        face_info = sorted(
            face_info_list,
            key=lambda x: (x["bbox"][2] - x["bbox"][0]) * (x["bbox"][3] - x["bbox"][1]),
        )[-1]
        
        # Detect gender from face info if available (for identity preservation)
        detected_gender = None
        if "gender" in face_info:
            detected_gender = face_info["gender"]
            print(f"  Detected gender: {'Female' if detected_gender == 0 else 'Male' if detected_gender == 1 else 'Unknown'}")
        elif hasattr(face_info, "get") and face_info.get("gender") is not None:
            detected_gender = face_info.get("gender")
            print(f"  Detected gender: {'Female' if detected_gender == 0 else 'Male' if detected_gender == 1 else 'Unknown'}")

        # Configure scheduler
        print("Step 5/8: Configuring scheduler...")
        scheduler_class_name = scheduler.split("-")[0]
        add_kwargs = {}
        if len(scheduler.split("-")) > 1:
            add_kwargs["use_karras_sigmas"] = True
        if len(scheduler.split("-")) > 2:
            add_kwargs["algorithm_type"] = "sde-dpmsolver++"
        scheduler_cls = getattr(diffusers, scheduler_class_name)
        pipe.scheduler = scheduler_cls.from_config(pipe.scheduler.config, **add_kwargs)
        print(f"  ✓ Scheduler: {scheduler_class_name}")

        # Apply style and process prompts (AFTER face detection so we can use gender info)
        if not prompt:
            prompt = "a person"
        
        # Add explicit gender to prompt if detected (InstantID works better with explicit gender)
        if detected_gender is not None:
            prompt_lower = prompt.lower()
            # Only add gender if not already in prompt
            if "man" not in prompt_lower and "woman" not in prompt_lower and "male" not in prompt_lower and "female" not in prompt_lower and "person" not in prompt_lower:
                if detected_gender == 0:  # Female
                    prompt = f"a woman, {prompt}"
                    print(f"  ✓ Added 'a woman' to prompt for gender preservation")
                elif detected_gender == 1:  # Male
                    prompt = f"a man, {prompt}"
                    print(f"  ✓ Added 'a man' to prompt for gender preservation")
            elif "person" in prompt_lower:
                # Replace "person" with specific gender
                if detected_gender == 0:  # Female
                    prompt = prompt.replace("person", "woman").replace("Person", "Woman")
                    print(f"  ✓ Replaced 'person' with 'woman' in prompt")
                elif detected_gender == 1:  # Male
                    prompt = prompt.replace("person", "man").replace("Person", "Man")
                    print(f"  ✓ Replaced 'person' with 'man' in prompt")
        
        # Warn if prompt contains physical feature descriptions that might override identity
        physical_keywords = ["hair", "blonde", "brown hair", "black hair", "red hair", "beard", "mustache", 
                            "wearing", "shirt", "jacket", "suit", "blazer", "tie", "glasses"]
        prompt_lower = prompt.lower()
        if any(keyword in prompt_lower for keyword in physical_keywords):
            print("  ⚠ Warning: Prompt contains physical feature descriptions. These may override face identity.")
            print("  💡 Tip: Focus on style/setting only (e.g., 'professional headshot, studio lighting') for better identity preservation.")
        
        # Add gender preservation to negative prompt if gender was detected
        gender_negative_terms = "wrong gender, gender swap, different person, different face, face swap, identity change, different identity"
        if detected_gender is not None:
            # Add opposite gender terms to negative prompt
            if detected_gender == 0:  # Female
                gender_negative_terms += ", man, male, masculine, boy"
            elif detected_gender == 1:  # Male
                gender_negative_terms += ", woman, female, feminine, girl"
            print(f"  ✓ Gender preservation enabled in negative prompt")
        
        # Add gender preservation terms to negative prompt
        if gender_negative_terms not in negative_prompt:
            negative_prompt = f"{negative_prompt}, {gender_negative_terms}" if negative_prompt else gender_negative_terms
        
        prompt, negative_prompt = apply_style(style_name, prompt, negative_prompt)
        print(f"  ✓ Style applied: {style_name}")
        
        # Extract face embedding and ensure it's a proper tensor
        face_emb_raw = face_info["embedding"]
        print(f"  Face embedding raw type: {type(face_emb_raw)}")
        
        # Convert to torch tensor if it's numpy
        if isinstance(face_emb_raw, np.ndarray):
            face_emb = torch.from_numpy(face_emb_raw).to(device).to(dtype)
        elif isinstance(face_emb_raw, torch.Tensor):
            face_emb = face_emb_raw.to(device).to(dtype)
        else:
            # Fallback: try to convert
            face_emb = torch.tensor(face_emb_raw, device=device, dtype=dtype)
        
        # Ensure proper shape for InstantID (should be [1, 512] or [512])
        if len(face_emb.shape) == 1:
            face_emb = face_emb.unsqueeze(0)  # Add batch dimension: [1, 512]
        elif len(face_emb.shape) == 2 and face_emb.shape[0] > 1:
            # If multiple faces, take the first one
            face_emb = face_emb[0:1]
        
        # DO NOT normalize - InstantID expects raw embeddings from InsightFace
        # Normalization breaks the identity preservation
        
        print(f"  Face embedding final shape: {face_emb.shape}, dtype: {face_emb.dtype}, device: {face_emb.device}")
        print(f"  Face embedding range: [{face_emb.min().item():.4f}, {face_emb.max().item():.4f}]")
        
        face_kps = draw_kps(convert_from_cv2_to_image(face_image_cv2), face_info["kps"])
        print(f"  Face keypoints image size: {face_kps.size}")
        img_controlnet = face_image
        print("  ✓ Face features extracted")
        import sys
        sys.stdout.flush()

        # Create control mask if requested
        print("Step 5/8: Preparing control images...")
        if enhance_face_region:
            control_mask = np.zeros([height, width, 3])
            x1, y1, x2, y2 = face_info["bbox"]
            x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
            control_mask[y1:y2, x1:x2] = 255
            control_mask = Image.fromarray(control_mask.astype(np.uint8))
        else:
            control_mask = None

        # Configure ControlNet
        print("Step 6/8: Configuring ControlNet...")
        if len(controlnet_selection) > 0:
            controlnet_scales = {
                "canny": canny_strength,
                "depth": depth_strength,
            }
            
            # Generate control images with progress updates first
            control_images = [face_kps]
            successful_controls = []
            
            for s in controlnet_selection:
                print(f"  Generating {s} control image...")
                try:
                    if s == "depth":
                        # Depth generation can be slow, show progress
                        print("    (This may take 10-20 seconds on CPU...)")
                    control_img = controlnet_map_fn[s](img_controlnet).resize((width, height))
                    control_images.append(control_img)
                    successful_controls.append(s)
                    print(f"  ✓ {s} control image generated")
                except Exception as e:
                    print(f"  ⚠ Warning: Failed to generate {s} control image: {e}")
                    print(f"    Continuing without {s} control...")
            
            # Configure ControlNet based on successful controls
            if len(successful_controls) > 0:
                pipe.controlnet = MultiControlNetModel(
                    [controlnet_identitynet]
                    + [controlnet_map[s] for s in successful_controls]
                )
                control_scales = [float(identitynet_strength_ratio)] + [
                    controlnet_scales[s] for s in successful_controls
                ]
                print(f"  ✓ ControlNet configured with: {successful_controls}")
            else:
                # Fallback to identity only if all controls failed
                pipe.controlnet = controlnet_identitynet
                control_scales = float(identitynet_strength_ratio)
                control_images = face_kps
                print("  ✓ ControlNet configured: identity only (control generation failed)")
        else:
            pipe.controlnet = controlnet_identitynet
            control_scales = float(identitynet_strength_ratio)
            control_images = face_kps
            print("  ✓ ControlNet configured: identity only")

        # Adjust steps for LCM if enabled
        if enable_LCM:
            num_steps = max(5, min(num_steps, 10))
            guidance_scale = max(1.0, min(guidance_scale, 2.0))
            print("  ✓ Fast generation mode enabled")

        generator = torch.Generator(device=device).manual_seed(seed)
        print(f"  Generator created on {device}, seed: {seed}")

        pipe.set_ip_adapter_scale(adapter_strength_ratio)
        print(f"  IP-Adapter scale set to: {adapter_strength_ratio}")
        import sys
        sys.stdout.flush()
        
        # Generate image
        print("Step 7/8: Generating image (this may take 30-60 seconds)...")
        print(f"  Steps: {num_steps}, Guidance: {guidance_scale}, Seed: {seed}")
        print(f"  Device: {device}, Image size: {width}x{height}")
        print(f"  Control images type: {type(control_images)}, length: {len(control_images) if isinstance(control_images, list) else 'single'}")
        print("  Starting diffusion process...")
        import sys
        sys.stdout.flush()  # Force output
        
        # Create callback for progress updates
        step_count = [0]  # Use list to allow modification in nested function
        
        def progress_callback(step, timestep, latents):
            step_count[0] = step + 1
            print(f"  Progress: Step {step_count[0]}/{num_steps} ({(step_count[0]/num_steps)*100:.1f}%)")
            sys.stdout.flush()
            if progress:
                try:
                    progress(step / num_steps, desc=f"Generating... {step}/{num_steps}")
                except:
                    pass
        
        print("  Calling pipeline...")
        sys.stdout.flush()
        
        try:
            # Ensure all inputs are on correct device
            print("  Preparing inputs...")
            print(f"    Prompt length: {len(prompt)} chars")
            print(f"    Negative prompt length: {len(negative_prompt)} chars")
            print(f"    Face embedding: shape={face_emb.shape}, dtype={face_emb.dtype}, device={face_emb.device}")
            print(f"    Face embedding has data: {face_emb.numel() > 0}")
            if isinstance(control_images, list):
                print(f"    Control images: {len(control_images)} images")
                for i, img in enumerate(control_images):
                    print(f"      Image {i}: {type(img)}, size: {img.size if hasattr(img, 'size') else 'N/A'}")
            else:
                print(f"    Control image: {type(control_images)}, size: {control_images.size if hasattr(control_images, 'size') else 'N/A'}")
            sys.stdout.flush()
            
            # Verify face_emb is not a meta tensor
            if hasattr(face_emb, 'is_meta') and face_emb.is_meta:
                print("  ✗ ERROR: Face embedding is a meta tensor!")
                raise RuntimeError("Face embedding is a meta tensor. This usually means the model wasn't loaded correctly.")
            
            # Ensure face_emb has actual data
            if face_emb.numel() == 0:
                print("  ✗ ERROR: Face embedding is empty!")
                raise RuntimeError("Face embedding is empty. Face detection may have failed.")
            
            # Check if pipeline is ready
            print("  Checking pipeline state...")
            print(f"    Pipeline device: {next(pipe.unet.parameters()).device}")
            print(f"    Pipeline dtype: {next(pipe.unet.parameters()).dtype}")
            sys.stdout.flush()
            
            # Force garbage collection before generation
            import gc
            if device == "cuda":
                torch.cuda.empty_cache()
            gc.collect()
            print("  Memory cleared, starting generation...")
            sys.stdout.flush()
            
            # Call pipeline with explicit error handling
            print("  Entering pipeline call (this may take a moment to start)...")
            print("  NOTE: First step may take 10-30 seconds for encoding/preprocessing")
            sys.stdout.flush()
            
            # Try to enable attention slicing for memory efficiency
            try:
                if hasattr(pipe, 'enable_attention_slicing'):
                    pipe.enable_attention_slicing(slice_size="max")
                    print("  ✓ Attention slicing enabled")
            except:
                pass
            
            # Try to enable CPU offload if on CPU
            if device == "cpu":
                try:
                    if hasattr(pipe, 'enable_sequential_cpu_offload'):
                        pipe.enable_sequential_cpu_offload()
                        print("  ✓ Sequential CPU offload enabled")
                except:
                    pass
            
            sys.stdout.flush()
            
            # Add a print right before the actual call
            print("  Starting pipeline inference NOW...")
            sys.stdout.flush()
            
            # Pass IP-Adapter scale explicitly to ensure it's used
            images = pipe(
                prompt=prompt,
                negative_prompt=negative_prompt,
                image_embeds=face_emb,
                image=control_images,
                control_mask=control_mask,
                controlnet_conditioning_scale=control_scales,
                num_inference_steps=num_steps,
                guidance_scale=guidance_scale,
                height=height,
                width=width,
                generator=generator,
                callback=progress_callback,
                callback_steps=1,  # Update every step
                ip_adapter_scale=adapter_strength_ratio,  # Explicitly pass IP-Adapter scale
            ).images
            
            print(f"  ✓ Pipeline completed, generated {len(images)} image(s)")
            sys.stdout.flush()
        except Exception as pipe_error:
            print(f"  ✗ Pipeline error: {pipe_error}")
            import traceback
            traceback.print_exc()
            sys.stdout.flush()
            raise

        final_image = images[0]
        print("Step 8/8: Saving image...")
        save_as_png(final_image)
        print("✓ Image generation complete!")
        print("=" * 60)
        
        # Clean up GPU memory before returning (important for ZeroGPU)
        if device == "cuda":
            try:
                torch.cuda.empty_cache()
                print("  GPU memory cleared")
            except Exception as e:
                print(f"  Warning: Could not clear GPU cache: {e}")
        
        # Force garbage collection
        gc.collect()
        
        # Ensure stdout is flushed before returning (helps with ZeroGPU)
        sys.stdout.flush()
        
        return final_image

    except gr.Error:
        raise
    except Exception as e:
        print(f"\n✗ Error during generation: {e}")
        import traceback
        traceback.print_exc()
        raise gr.Error(f"An error occurred during generation: {str(e)}")


# ============ CSS STYLING ============
css = """
/* Main container styling */
.main-container {
    max-width: 1400px;
    margin: 0 auto;
    padding: 20px;
}

/* Hero section */
.hero-title {
    font-size: 2.5em;
    font-weight: 700;
    margin-bottom: 10px;
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    background-clip: text;
}

.hero-subtitle {
    font-size: 1.1em;
    color: #666;
    margin-bottom: 30px;
}

/* Control cards */
.control-card {
    background: #f8f9fa;
    border-radius: 12px;
    padding: 20px;
    margin-bottom: 20px;
    box-shadow: 0 2px 4px rgba(0,0,0,0.1);
}

.control-header {
    display: flex;
    align-items: center;
    margin-bottom: 15px;
}

.control-icon {
    font-size: 1.5em;
    margin-right: 10px;
}

.control-title {
    font-size: 1.2em;
    font-weight: 600;
    margin: 0;
}

/* Upload area */
.upload-area {
    border: 2px dashed #667eea;
    border-radius: 8px;
    padding: 20px;
    text-align: center;
}

/* Tips card */
.tips-card {
    background: #fff3cd;
    border-left: 4px solid #ffc107;
    border-radius: 8px;
    padding: 15px;
    margin-bottom: 20px;
}

.tips-header {
    display: flex;
    align-items: center;
    margin-bottom: 10px;
}

.tips-icon {
    font-size: 1.3em;
    margin-right: 8px;
}

/* Result card */
.result-card {
    background: #f8f9fa;
    border-radius: 12px;
    padding: 20px;
    box-shadow: 0 2px 4px rgba(0,0,0,0.1);
}

.result-header {
    margin-bottom: 20px;
}

.result-title {
    font-size: 1.5em;
    font-weight: 600;
    margin-bottom: 5px;
}

.result-subtitle {
    color: #666;
    font-size: 0.95em;
}

/* Image container */
.image-container {
    border-radius: 8px;
    overflow: hidden;
    box-shadow: 0 4px 8px rgba(0,0,0,0.1);
}

/* Success banner */
.success-banner {
    background: #d4edda;
    border: 1px solid #c3e6cb;
    border-radius: 8px;
    padding: 15px;
    margin-top: 15px;
}

/* Primary button */
.btn-primary {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    color: white;
    border: none;
    font-weight: 600;
    padding: 12px 30px;
}

.btn-primary:hover {
    opacity: 0.9;
    transform: translateY(-1px);
    box-shadow: 0 4px 8px rgba(0,0,0,0.2);
}
"""


# ============ UI / GRADIO ============
def show_success():
    """Show success message after generation."""
    return gr.update(
        value="""
        <div class="success-banner">
            <h4 style="margin: 0 0 8px 0;">✅ Success! Your Professional Headshot is Ready</h4>
            <p style="margin: 0; opacity: 0.9;">Download your high-quality PNG file for LinkedIn, professional profiles, or portfolios.</p>
        </div>
        """,
    )


with gr.Blocks() as demo:
    # Add custom CSS via HTML component (compatible with Gradio 4.44.0+)
    gr.HTML(f"<style>{css}</style>")
    
    with gr.Column(elem_classes="main-container"):
        with gr.Column(elem_classes="hero-section"):
            gr.HTML(
                """
                <div style="position: relative; z-index: 2;">
                    <h1 class="hero-title">🎯 Pro AI Headshot Generator</h1>
                    <p class="hero-subtitle">Transform any selfie into professional headshots in seconds. Perfect for LinkedIn, corporate profiles, and professional portfolios.</p>
                </div>
                """
            )

        with gr.Row(equal_height=True):
            with gr.Column(scale=1, min_width=400):
                with gr.Column(elem_classes="control-card"):
                    gr.HTML(
                        """
                        <div class="control-header">
                            <div class="control-icon">📸</div>
                            <h3 class="control-title">Upload Your Photo</h3>
                        </div>
                        """
                    )
                    gr.HTML(
                        """
                        <p style="color: var(--text-secondary); margin-bottom: 20px; font-size: 0.95em;">
                            For best results, use a clear, well-lit photo where your face is clearly visible.
                        </p>
                        """
                    )
                    face_file = gr.Image(
                        label="",
                        type="filepath",
                        height=200,
                        show_label=False,
                        elem_classes="upload-area",
                    )

                with gr.Column(elem_classes="control-card"):
                    gr.HTML(
                        """
                        <div class="control-header">
                            <div class="control-icon">✍️</div>
                            <h3 class="control-title">Describe Your Look</h3>
                        </div>
                        """
                    )
                    prompt = gr.Textbox(
                        label="",
                        placeholder="Describe the style and setting (avoid describing physical features)...",
                        value="professional headshot, business portrait, soft natural lighting, high quality",
                        show_label=False,
                        lines=3,
                    )
                    gr.HTML(
                        """
                        <div style="font-size: 0.85em; color: var(--text-secondary); margin-top: 8px;">
                            💡 Tips: Describe style/setting only. Don't mention hair color, clothing, or physical features - your face will be preserved automatically.
                            <br>Examples: "professional business headshot", "corporate portrait", "studio lighting"
                        </div>
                        """
                    )

                with gr.Column(elem_classes="control-card"):
                    gr.HTML(
                        """
                        <div class="control-header">
                            <div class="control-icon">🎨</div>
                            <h3 class="control-title">Style Options</h3>
                        </div>
                        """
                    )
                    style = gr.Dropdown(
                        label="Style Theme",
                        choices=["No Style"] + STYLE_NAMES,
                        value="No Style",
                        info="'No Style' recommended for natural professional results",
                    )

                with gr.Column(elem_classes="control-card"):
                    gr.HTML(
                        """
                        <div class="control-header">
                            <div class="control-icon">⚙️</div>
                            <h3 class="control-title">Quality Settings</h3>
                        </div>
                        """
                    )

                    identitynet_strength_ratio = gr.Slider(
                        label="Face Similarity",
                        minimum=0.5,
                        maximum=1.5,
                        step=0.05,
                        value=1.2,
                        info="How closely the headshot resembles your photo (higher = more similar, recommended: 1.0-1.2)",
                    )

                    adapter_strength_ratio = gr.Slider(
                        label="Face Detail Strength",
                        minimum=0.3,
                        maximum=1.5,
                        step=0.05,
                        value=1.0,
                        info="Strength of face features preservation (higher = better identity match, recommended: 0.8-1.0)",
                    )

                    enable_LCM = gr.Checkbox(
                        label="Enable Fast Generation Mode",
                        value=False,
                        info="Faster results with slightly lower quality",
                    )

                with gr.Column(elem_classes="tips-card"):
                    gr.HTML(
                        """
                        <div class="tips-header">
                            <div class="tips-icon">💡</div>
                            <h4 style="margin: 0; color: #92400e;">Pro Tips for Best Results</h4>
                        </div>
                        <ul style="margin: 0; color: #92400e; font-size: 0.9em;">
                            <li>Use clear, well-lit face photos</li>
                            <li>Face should be visible and not too small</li>
                            <li>Avoid blurry or dark images</li>
                            <li>Single person in photo works best</li>
                        </ul>
                        """
                    )

                submit = gr.Button(
                    "✨ Generate Professional Headshot",
                    variant="primary",
                    size="lg",
                    elem_classes="btn-primary",
                    scale=1,
                )

            with gr.Column(scale=1, min_width=500):
                with gr.Column(elem_classes="result-card"):
                    gr.HTML(
                        """
                        <div class="result-header">
                            <h2 class="result-title">Your Professional Headshot</h2>
                            <p class="result-subtitle">Your AI-generated headshot will appear here. Download as high-quality PNG for professional use.</p>
                        </div>
                        """
                    )

                    gallery = gr.Image(
                        label="Output",
                        height=400,
                        show_label=False,
                        type="pil",
                        elem_classes="image-container",
                    )

                    success_msg = gr.HTML(
                        """
                        <div class="success-banner" style="display: none;">
                            <h4 style="margin: 0 0 8px 0;">✅ Success! Your Professional Headshot is Ready</h4>
                            <p style="margin: 0; opacity: 0.9;">Download your high-quality PNG file for LinkedIn, professional profiles, or portfolios.</p>
                        </div>
                        """
                    )

                    progress_info = gr.HTML(
                        """
                        <div class="progress-container">
                            <div style="font-size: 0.9em; color: var(--text-secondary);">
                                ⏱️ Generation takes 20-30 seconds
                            </div>
                        </div>
                        """
                    )

    # Hidden advanced settings
    negative_prompt = gr.Textbox(
        value="(lowres, low quality, worst quality:1.2), (text:1.2), watermark, (frame:1.2), deformed, ugly, deformed eyes, blur, out of focus, blurry, deformed cat, deformed, photo, anthropomorphic cat, monochrome, pet collar, gun, weapon, blue, 3d, drones, drone, buildings in background, green, wrong gender, gender swap, different person, different face, face swap, identity change, different identity",
        visible=False,
    )
    num_steps = gr.Slider(
        minimum=5,
        maximum=100,
        step=1,
        value=DEFAULT_NUM_STEPS,
        label="Number of steps",
        visible=False,
    )
    guidance_scale = gr.Slider(
        minimum=0.1,
        maximum=20.0,
        step=0.1,
        value=DEFAULT_GUIDANCE_SCALE,
        label="Guidance scale",
        visible=False,
    )
    seed = gr.Slider(
        minimum=-1,
        maximum=MAX_SEED,
        step=1,
        value=-1,
        label="Seed (-1 for random)",
        visible=False,
    )
    scheduler = gr.Dropdown(
        value="EulerDiscreteScheduler",
        choices=[
            "EulerDiscreteScheduler",
            "EulerAncestralDiscreteScheduler",
            "DPMSolverMultistepScheduler",
        ],
        visible=False,
    )
    randomize_seed = gr.Checkbox(value=True, visible=False)
    enhance_face_region = gr.Checkbox(value=True, visible=False)
    controlnet_selection = gr.CheckboxGroup(
        choices=["canny", "depth"], value=[], label="Controlnet", visible=False  # Changed default to empty - depth can be slow
    )
    canny_strength = gr.Slider(
        minimum=0,
        maximum=1.5,
        step=0.01,
        value=0.4,
        label="Canny strength",
        visible=False,
    )
    depth_strength = gr.Slider(
        minimum=0,
        maximum=1.5,
        step=0.01,
        value=0.4,
        label="Depth strength",
        visible=False,
    )

    submit.click(
        fn=generate_image,
        inputs=[
            face_file,
            prompt,
            negative_prompt,
            style,
            num_steps,
            identitynet_strength_ratio,
            adapter_strength_ratio,
            canny_strength,
            depth_strength,
            controlnet_selection,
            guidance_scale,
            seed,
            scheduler,
            enable_LCM,
            enhance_face_region,
        ],
        outputs=[gallery],
    ).then(fn=show_success, outputs=success_msg)

    enable_LCM.input(
        fn=toggle_lcm_ui,
        inputs=[enable_LCM],
        outputs=[num_steps, guidance_scale],
        queue=False,
    )

    # Cleanup on startup
    cleanup_old_files()


if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description="Pro AI Headshot Generator")
    parser.add_argument("--share", action="store_true", help="Create a public Gradio link")
    parser.add_argument("--server-name", type=str, default="127.0.0.1", help="Server name (default: 127.0.0.1)")
    parser.add_argument("--server-port", type=int, default=7860, help="Server port (default: 7860)")
    parser.add_argument("--no-queue", action="store_true", help="Disable queueing")
    args = parser.parse_args()
    
    demo.queue(api_open=False, max_size=3 if not args.no_queue else None)
    
    # Launch the application
    try:
        print("\n" + "=" * 60)
        print("🚀 Launching Gradio interface...")
        print("=" * 60)
        
        # For Hugging Face Spaces, use 0.0.0.0 to allow external access
        server_name = "0.0.0.0" if os.getenv("SPACE_ID") else args.server_name
        
        print(f"  Server: {server_name}:{args.server_port}")
        print(f"  Queue enabled: {not args.no_queue}")
        print("  Interface should be available shortly...")
        print("=" * 60)
        
        demo.launch(
            share=args.share,
            server_name=server_name,
            server_port=args.server_port,
            show_error=True,
            # Note: theme parameter removed for Gradio 4.44.0+ compatibility
            # CSS is injected via HTML component instead
        )
        
        print("\n✅ Gradio interface launched successfully!")
    except Exception as e:
        print(f"\n✗ Error launching Gradio interface: {e}")
        import traceback
        traceback.print_exc()
        raise