initial demo implementation

app.py

@@ -1,7 +1,532 @@
+import os
+import sys
+os.environ["OMP_NUM_THREADS"] = "1"
+os.environ["MKL_NUM_THREADS"] = "1"
+os.environ["MKL_THREADING_LAYER"] = "GNU"
+
+# ---------------------------------------------------------------------------
+# Configuration — set CHECKPOINT_REPO as a HuggingFace Space secret to load
+# fine-tuned models. If left empty, the demo uses base Arc2Face with a raw
+# WavLM x-vector encoder (useful for testing that the Space works).
+# ---------------------------------------------------------------------------
+CHECKPOINT_REPO = os.environ.get("CHECKPOINT_REPO", "")
+ENCODER_FILENAME = os.environ.get("ENCODER_FILENAME", "speaker_encoder.pt")
+ARC2FACE_REPO = "FoivosPar/Arc2Face"
+BASE_MODEL = "stable-diffusion-v1-5/stable-diffusion-v1-5"
+SKIP_LORA = not bool(CHECKPOINT_REPO)
+SKIP_SPEAKER_ENCODER = not bool(CHECKPOINT_REPO)
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchaudio
+from PIL import Image
+from diffusers import StableDiffusionPipeline, UNet2DConditionModel, DPMSolverMultistepScheduler
+from huggingface_hub import snapshot_download, hf_hub_download
 import gradio as gr
 
-def greet(name):
-    return "Hello " + name + "!!"
+from external.arc2face import CLIPTextModelWrapper, project_face_embs
+from core.models.encoder.speech_face_encoder import SpeechFaceXVectorEncoder
+
+# ---------------------------------------------------------------------------
+# Globals populated at startup
+# ---------------------------------------------------------------------------
+pipeline = None
+speaker_encoder = None
+facenet_model = None
+facenet_classify_model = None
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+
+# ---------------------------------------------------------------------------
+# PEFT-compatible attention processors (inlined from core/factories/lora_factory.py)
+# These fix "Linear.forward() takes 2 positional arguments but 3 were given"
+# when using LoRA-wrapped UNet attention layers.
+# ---------------------------------------------------------------------------
+
+class PeftCompatibleAttnProcessor:
+    def __call__(
+        self,
+        attn,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states=None,
+        attention_mask=None,
+        temb=None,
+        *args,
+        **kwargs,
+    ) -> torch.Tensor:
+        residual = hidden_states
+
+        if attn.spatial_norm is not None:
+            hidden_states = attn.spatial_norm(hidden_states, temb)
+
+        input_ndim = hidden_states.ndim
+
+        if input_ndim == 4:
+            batch_size, channel, height, width = hidden_states.shape
+            hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+
+        batch_size, sequence_length, _ = (
+            hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+        )
+        attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+
+        if attn.group_norm is not None:
+            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
+
+        query = attn.to_q(hidden_states)
+
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states
+        elif attn.norm_cross:
+            encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
+
+        key = attn.to_k(encoder_hidden_states)
+        value = attn.to_v(encoder_hidden_states)
+
+        query = attn.head_to_batch_dim(query)
+        key = attn.head_to_batch_dim(key)
+        value = attn.head_to_batch_dim(value)
+
+        attention_probs = attn.get_attention_scores(query, key, attention_mask)
+        hidden_states = torch.bmm(attention_probs, value)
+        hidden_states = attn.batch_to_head_dim(hidden_states)
+
+        hidden_states = attn.to_out[0](hidden_states)
+        hidden_states = attn.to_out[1](hidden_states)
+
+        if input_ndim == 4:
+            hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
+
+        if attn.residual_connection:
+            hidden_states = hidden_states + residual
+
+        hidden_states = hidden_states / attn.rescale_output_factor
+        return hidden_states
+
+
+class PeftCompatibleAttnProcessor2_0:
+    def __init__(self):
+        if not hasattr(torch.nn.functional, "scaled_dot_product_attention"):
+            raise ImportError("PeftCompatibleAttnProcessor2_0 requires PyTorch 2.0+.")
+
+    def __call__(
+        self,
+        attn,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states=None,
+        attention_mask=None,
+        temb=None,
+        *args,
+        **kwargs,
+    ) -> torch.Tensor:
+        residual = hidden_states
+
+        if attn.spatial_norm is not None:
+            hidden_states = attn.spatial_norm(hidden_states, temb)
+
+        input_ndim = hidden_states.ndim
+
+        if input_ndim == 4:
+            batch_size, channel, height, width = hidden_states.shape
+            hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+
+        batch_size, sequence_length, _ = (
+            hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+        )
+
+        if attention_mask is not None:
+            attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+            attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
+
+        if attn.group_norm is not None:
+            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
+
+        query = attn.to_q(hidden_states)
+
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states
+        elif attn.norm_cross:
+            encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
+
+        key = attn.to_k(encoder_hidden_states)
+        value = attn.to_v(encoder_hidden_states)
+
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        hidden_states = torch.nn.functional.scaled_dot_product_attention(
+            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+        )
+
+        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+        hidden_states = hidden_states.to(query.dtype)
+
+        hidden_states = attn.to_out[0](hidden_states)
+        hidden_states = attn.to_out[1](hidden_states)
+
+        if input_ndim == 4:
+            hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
+
+        if attn.residual_connection:
+            hidden_states = hidden_states + residual
+
+        hidden_states = hidden_states / attn.rescale_output_factor
+        return hidden_states
+
+
+def _set_attn_processor_for_lora(unet: nn.Module) -> None:
+    try:
+        attn_procs = {}
+        for name in unet.attn_processors.keys():
+            if hasattr(torch.nn.functional, 'scaled_dot_product_attention'):
+                attn_procs[name] = PeftCompatibleAttnProcessor2_0()
+            else:
+                attn_procs[name] = PeftCompatibleAttnProcessor()
+        unet.set_attn_processor(attn_procs)
+        print("  Set PEFT-compatible attention processors")
+    except Exception as e:
+        print(f"  Warning: Could not set attention processors for LoRA: {e}")
+
+
+# ---------------------------------------------------------------------------
+# Utilities
+# ---------------------------------------------------------------------------
+
+def load_and_process_audio(audio_file: str, dev: str, max_seconds: float = 6.0):
+    try:
+        waveform, sample_rate = torchaudio.load(audio_file)
+    except Exception:
+        import soundfile as sf
+        data, sample_rate = sf.read(audio_file, always_2d=True)
+        waveform = torch.from_numpy(data.T.astype(np.float32))
+    if sample_rate != 16000:
+        resampler = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=16000)
+        waveform = resampler(waveform)
+    if waveform.shape[0] > 1:
+        waveform = waveform.mean(dim=0, keepdim=True)
+    max_samples = int(max_seconds * 16000)
+    if waveform.shape[1] > max_samples:
+        waveform = waveform[:, :max_samples]
+    elif waveform.shape[1] < max_samples:
+        waveform = F.pad(waveform, (0, max_samples - waveform.shape[1]))
+    return waveform.squeeze(0).unsqueeze(0).to(dev)
+
+
+def is_lora_checkpoint(checkpoint_path: str, subfolder: str) -> bool:
+    return os.path.exists(os.path.join(checkpoint_path, subfolder, "adapter_config.json"))
+
+
+def resolve_checkpoint_path(checkpoint_path: str) -> str:
+    checkpoint_path = os.path.expanduser(checkpoint_path)
+    if not os.path.exists(checkpoint_path):
+        raise FileNotFoundError(f"Checkpoint path does not exist: {checkpoint_path}")
+    expected_subs = {"encoder", "unet"}
+    if os.path.isdir(checkpoint_path):
+        children = set(os.listdir(checkpoint_path))
+        if expected_subs.issubset(children):
+            return checkpoint_path
+        ckpts = [d for d in os.listdir(checkpoint_path)
+                 if d.startswith("checkpoint-") and os.path.isdir(os.path.join(checkpoint_path, d))]
+        if not ckpts:
+            return checkpoint_path
+
+        def ckpt_num(name):
+            try:
+                return int(name.split("checkpoint-")[-1])
+            except Exception:
+                return -1
+        return os.path.join(checkpoint_path, sorted(ckpts, key=ckpt_num)[-1])
+    return checkpoint_path
+
+
+# ---------------------------------------------------------------------------
+# LoRA checkpoint loading
+# ---------------------------------------------------------------------------
+
+def load_encoder_with_lora(checkpoint_path: str):
+    encoder_path = os.path.join(checkpoint_path, "lora", "encoder")
+    if is_lora_checkpoint(checkpoint_path, os.path.join("lora", "encoder")):
+        from peft import PeftModel
+        base_encoder = CLIPTextModelWrapper.from_pretrained(ARC2FACE_REPO, subfolder='encoder')
+        encoder = PeftModel.from_pretrained(base_encoder, encoder_path)
+        encoder = encoder.merge_and_unload()
+        encoder.forward = base_encoder.forward
+        return encoder
+    return CLIPTextModelWrapper.from_pretrained(checkpoint_path, subfolder="encoder")
+
+
+def load_unet_with_lora(checkpoint_path: str):
+    unet_path = os.path.join(checkpoint_path, "lora", "unet")
+    if is_lora_checkpoint(checkpoint_path, os.path.join("lora", "unet")):
+        from peft import PeftModel
+        base_unet = UNet2DConditionModel.from_pretrained(ARC2FACE_REPO, subfolder='arc2face')
+        unet = PeftModel.from_pretrained(base_unet, unet_path)
+        unet = unet.merge_and_unload()
+        unet.forward = base_unet.forward
+        _set_attn_processor_for_lora(unet)
+        return unet
+    return UNet2DConditionModel.from_pretrained(checkpoint_path, subfolder="unet")
+
+
+# ---------------------------------------------------------------------------
+# Raw WavLM encoder (fallback when no fine-tuned checkpoint is provided)
+# ---------------------------------------------------------------------------
+
+class RawWavLMEncoder:
+    def __init__(self, pretrained_path: str, dev: str):
+        from transformers import WavLMForXVector
+        self.wavlm_xvector = WavLMForXVector.from_pretrained(pretrained_path).to(dev)
+        self.wavlm_xvector.eval()
+
+    def __call__(self, waveform, normalize=True, apply_shared_projection=False):
+        emb = self.wavlm_xvector(input_values=waveform, return_dict=True).embeddings
+        if normalize:
+            emb = F.normalize(emb, p=2, dim=1)
+        return emb
+
+    def eval(self):
+        self.wavlm_xvector.eval()
+        return self
+
+    def to(self, dev):
+        self.wavlm_xvector = self.wavlm_xvector.to(dev)
+        return self
+
+
+# ---------------------------------------------------------------------------
+# FaceNet best-sample selection
+# ---------------------------------------------------------------------------
+
+def _facenet_transform():
+    from torchvision import transforms
+    return transforms.Compose([
+        transforms.Resize((160, 160)),
+        transforms.ToTensor(),
+        transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),
+    ])
+
+
+def _extract_facenet_emb(img: Image.Image, model) -> torch.Tensor:
+    tensor = _facenet_transform()(img.convert("RGB")).unsqueeze(0)
+    with torch.no_grad():
+        emb = model(tensor)
+    return F.normalize(emb.squeeze(0), p=2, dim=0)
+
+
+def _extract_facenet_logits(img: Image.Image, model) -> torch.Tensor:
+    tensor = _facenet_transform()(img.convert("RGB")).unsqueeze(0)
+    with torch.no_grad():
+        logits = model(tensor)
+    return logits.squeeze(0)
+
+
+def select_best_image(images: list, method: str) -> Image.Image:
+    global facenet_model, facenet_classify_model
+
+    if method == "entropy":
+        logits_list = [_extract_facenet_logits(img, facenet_classify_model) for img in images]
+        logits_stack = torch.stack(logits_list)
+        probs = F.softmax(logits_stack, dim=1)
+        entropy = -(probs * (probs + 1e-10).log()).sum(dim=1)
+        best_idx = entropy.argmin().item()
+        print(f"[select_best:entropy] selected image {best_idx} (entropy={entropy[best_idx]:.3f})")
+
+    elif method == "pairwise":
+        embeddings = torch.stack([_extract_facenet_emb(img, facenet_model) for img in images])
+        sim_matrix = F.cosine_similarity(embeddings.unsqueeze(1), embeddings.unsqueeze(0), dim=2)
+        avg_sims = (sim_matrix.sum(dim=1) - 1) / (len(images) - 1)
+        best_idx = avg_sims.argmax().item()
+        print(f"[select_best:pairwise] selected image {best_idx} (avg_sim={avg_sims[best_idx]:.3f})")
+
+    else:  # mean
+        embeddings = torch.stack([_extract_facenet_emb(img, facenet_model) for img in images])
+        mean_emb = F.normalize(embeddings.mean(dim=0), p=2, dim=0)
+        sims = F.cosine_similarity(embeddings, mean_emb.unsqueeze(0))
+        best_idx = sims.argmax().item()
+        print(f"[select_best:mean] selected image {best_idx} (sim={sims[best_idx]:.3f})")
+
+    return images[best_idx]
+
+
+# ---------------------------------------------------------------------------
+# Generation
+# ---------------------------------------------------------------------------
+
+def generate(audio_path, num_samples, guidance_scale, num_inference_steps, base_seed, select_best, best_selection="pairwise"):
+    global pipeline, speaker_encoder, facenet_model, facenet_classify_model, device
+
+    if pipeline is None:
+        return None, "Model not loaded. Check Space configuration."
+    if audio_path is None:
+        return None, "Please provide an audio file."
+
+    try:
+        waveform = load_and_process_audio(audio_path, device, max_seconds=5.0)
+    except Exception as e:
+        return None, f"Audio loading failed: {e}"
+
+    with torch.no_grad():
+        speech_z = speaker_encoder(waveform, normalize=True, apply_shared_projection=False)
+        id_emb = speech_z.to(torch.float16)
+        id_emb_projected = project_face_embs(pipeline, id_emb)
+
+    images = []
+    for i in range(int(num_samples)):
+        seed = int(base_seed) + i
+        generator = torch.Generator(device=device).manual_seed(seed)
+        img = pipeline(
+            prompt_embeds=id_emb_projected,
+            num_inference_steps=int(num_inference_steps),
+            guidance_scale=float(guidance_scale),
+            num_images_per_prompt=1,
+            generator=generator,
+        ).images[0]
+        images.append(img)
+
+    if select_best:
+        model_ready = facenet_model is not None if best_selection in ("mean", "pairwise") else facenet_classify_model is not None
+        if model_ready:
+            best = select_best_image(images, best_selection)
+        else:
+            best = images[0]
+        return [best], ""
+
+    return images, ""
+
+
+# ---------------------------------------------------------------------------
+# Model loading
+# ---------------------------------------------------------------------------
+
+def load_models():
+    global pipeline, speaker_encoder, facenet_model, facenet_classify_model, device
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    print(f"Using device: {device}")
+
+    # Speaker encoder
+    print("Loading speaker encoder...")
+    if SKIP_SPEAKER_ENCODER:
+        speaker_encoder = RawWavLMEncoder("microsoft/wavlm-base-sv", device)
+        print("  Using raw WavLM x-vector encoder (no fine-tuned checkpoint)")
+    else:
+        enc = SpeechFaceXVectorEncoder(
+            pretrained_path="microsoft/wavlm-base-sv",
+            face_emb_dim=512,
+            dropout=0.0,
+            use_projection=True,
+            freeze_feature_encoder=True,
+        )
+        encoder_pt = hf_hub_download(CHECKPOINT_REPO, ENCODER_FILENAME)
+        ckpt = torch.load(encoder_pt, map_location=device, weights_only=False)
+        enc.load_state_dict(ckpt["model"], strict=False)
+        speaker_encoder = enc.to(device).eval()
+        print(f"  Loaded from {CHECKPOINT_REPO}/{ENCODER_FILENAME}")
+
+    # Diffusion pipeline
+    print("Loading diffusion pipeline...")
+    if SKIP_LORA:
+        encoder = CLIPTextModelWrapper.from_pretrained(ARC2FACE_REPO, subfolder='encoder', torch_dtype=torch.float16)
+        unet = UNet2DConditionModel.from_pretrained(ARC2FACE_REPO, subfolder='arc2face', torch_dtype=torch.float16)
+        print("  Using base Arc2Face (no LoRA)")
+    else:
+        checkpoint_dir = snapshot_download(CHECKPOINT_REPO)
+        checkpoint = resolve_checkpoint_path(checkpoint_dir)
+        print(f"  Checkpoint: {checkpoint}")
+        encoder = load_encoder_with_lora(checkpoint).to(dtype=torch.float16)
+        unet = load_unet_with_lora(checkpoint).to(dtype=torch.float16)
+
+    pipeline = StableDiffusionPipeline.from_pretrained(
+        BASE_MODEL,
+        text_encoder=encoder,
+        unet=unet,
+        torch_dtype=torch.float16,
+        safety_checker=None,
+    )
+    pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config)
+    pipeline = pipeline.to(device)
+    print("  Pipeline ready")
+
+    # FaceNet for best-sample selection
+    print("Loading FaceNet for best-sample selection...")
+    try:
+        from facenet_pytorch import InceptionResnetV1
+        facenet_model = InceptionResnetV1(pretrained='vggface2', classify=False).eval()
+        facenet_classify_model = InceptionResnetV1(pretrained='vggface2', classify=True).eval()
+        print("  FaceNet ready")
+    except Exception as e:
+        print(f"  FaceNet unavailable ({e}); select-best will fall back to first image")
+        facenet_model = None
+        facenet_classify_model = None
+
+
+# ---------------------------------------------------------------------------
+# Gradio UI
+# ---------------------------------------------------------------------------
+
+def build_demo():
+    facenet_available = facenet_model is not None and facenet_classify_model is not None
+
+    with gr.Blocks(title="Speech-to-Face Generation") as demo:
+        gr.Markdown("# Speech-to-Face Generation")
+        gr.Markdown("Upload or record a speech audio clip and generate face images conditioned on the speaker's voice.")
+
+        with gr.Row():
+            with gr.Column():
+                audio_input = gr.Audio(
+                    sources=["upload", "microphone"],
+                    type="filepath",
+                    label="Audio Input",
+                )
+                num_samples = gr.Slider(1, 16, value=4, step=1, label="Number of samples")
+                guidance_scale = gr.Slider(1.0, 10.0, value=4.5, step=0.5, label="Guidance scale")
+                num_steps = gr.Slider(10, 50, value=25, step=5, label="Inference steps")
+                base_seed = gr.Slider(0, 9999, value=42, step=1, label="Base seed")
+                select_best = gr.Checkbox(
+                    value=facenet_available,
+                    label="Select best image",
+                    interactive=facenet_available,
+                )
+                best_selection = gr.Radio(
+                    choices=["pairwise", "mean", "entropy"],
+                    value="pairwise",
+                    label="Selection method",
+                    info="pairwise: most consistent with others | mean: closest to average | entropy: most confident face prediction",
+                    interactive=facenet_available,
+                )
+                if not facenet_available:
+                    gr.Markdown("_FaceNet not available — select-best disabled._")
+                generate_btn = gr.Button("Generate", variant="primary")
+
+            with gr.Column():
+                gallery = gr.Gallery(label="Generated Images", columns=4, rows=1, object_fit="contain")
+                status = gr.Markdown(visible=False)
+
+        def _generate(audio, n, gs, steps, seed, best, selection):
+            imgs, msg = generate(audio, n, gs, steps, seed, best, selection)
+            visible = bool(msg)
+            return imgs, gr.update(value=msg, visible=visible)
+
+        generate_btn.click(
+            fn=_generate,
+            inputs=[audio_input, num_samples, guidance_scale, num_steps, base_seed, select_best, best_selection],
+            outputs=[gallery, status],
+        )
+
+    return demo
+
+
+# ---------------------------------------------------------------------------
+# Entry point
+# ---------------------------------------------------------------------------
+
+load_models()
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
+demo = build_demo()
 demo.launch()

core/models/encoder/speech_face_encoder.py

@@ -0,0 +1,438 @@
+"""
+Speech-Face Encoder Wrapper for WavLMForXVector
+
+Wraps HuggingFace's WavLMForXVector for speech-to-face embedding alignment task.
+Uses CLIP-style InfoNCE loss for cross-modal alignment.
+"""
+
+import logging
+import types
+from typing import Dict, List, Optional, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import WavLMForXVector, WavLMModel, WavLMConfig, AutoFeatureExtractor
+
+logger = logging.getLogger(__name__)
+
+
+class _TDNNLayer(nn.Module):
+    """TDNN layer matching HuggingFace WavLMTDNNLayer implementation exactly."""
+
+    def __init__(self, in_conv_dim: int, out_conv_dim: int, kernel_size: int, dilation: int):
+        super().__init__()
+        self.in_conv_dim = in_conv_dim
+        self.out_conv_dim = out_conv_dim
+        self.kernel_size = kernel_size
+        self.dilation = dilation
+        self.kernel = nn.Linear(in_conv_dim * kernel_size, out_conv_dim)
+        self.activation = nn.ReLU()
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        # hidden_states: (B, T, C)
+        hidden_states = hidden_states.unsqueeze(1)  # (B, 1, T, C)
+        hidden_states = F.unfold(
+            hidden_states,
+            (self.kernel_size, self.in_conv_dim),
+            stride=(1, self.in_conv_dim),
+            dilation=(self.dilation, 1),
+        )  # (B, kernel_size * in_conv_dim, T_out)
+        hidden_states = hidden_states.transpose(1, 2)  # (B, T_out, kernel_size * in_conv_dim)
+        hidden_states = self.kernel(hidden_states)
+        hidden_states = self.activation(hidden_states)
+        return hidden_states
+
+
+class WavLMBaseWithXVectorHead(nn.Module):
+    """
+    WavLM-base-plus backbone with a freshly-initialized x-vector head.
+    """
+
+    _TDNN_DIM = [512, 512, 512, 512, 1500]
+    _TDNN_KERNEL = [5, 3, 3, 1, 1]
+    _TDNN_DILATION = [1, 2, 3, 1, 1]
+    _XVECTOR_DIM = 512
+
+    def __init__(self, pretrained_path: str, use_weighted_layer_sum: bool = True):
+        super().__init__()
+
+        logger.info(f"Loading WavLMModel (base backbone) from {pretrained_path}")
+        self.wavlm = WavLMModel.from_pretrained(pretrained_path)
+        hidden_size = self.wavlm.config.hidden_size
+        num_layers = self.wavlm.config.num_hidden_layers
+
+        self._use_weighted_layer_sum = use_weighted_layer_sum
+        if use_weighted_layer_sum:
+            self.layer_weights = nn.Parameter(torch.ones(num_layers + 1))
+
+        self.projector = nn.Linear(hidden_size, self._TDNN_DIM[0])
+
+        in_dims = [self._TDNN_DIM[0]] + self._TDNN_DIM[:-1]
+        self.tdnn = nn.ModuleList([
+            _TDNNLayer(in_dims[i], self._TDNN_DIM[i], self._TDNN_KERNEL[i], self._TDNN_DILATION[i])
+            for i in range(len(self._TDNN_DIM))
+        ])
+
+        self.feature_extractor = nn.Linear(2 * self._TDNN_DIM[-1], self._XVECTOR_DIM)
+        self.classifier = nn.Linear(self._XVECTOR_DIM, self._XVECTOR_DIM)
+
+    @property
+    def config(self):
+        cfg = self.wavlm.config
+        cfg.xvector_output_dim = self._XVECTOR_DIM
+        cfg.use_weighted_layer_sum = self._use_weighted_layer_sum
+        return cfg
+
+    def freeze_base_model(self):
+        for param in self.wavlm.parameters():
+            param.requires_grad = False
+
+    def freeze_feature_encoder(self):
+        for param in self.wavlm.feature_extractor.parameters():
+            param.requires_grad = False
+
+    def _get_tdnn_output_lengths(self, feat_lengths: torch.Tensor) -> torch.Tensor:
+        for kernel, dilation in zip(self._TDNN_KERNEL, self._TDNN_DILATION):
+            feat_lengths = feat_lengths - dilation * (kernel - 1)
+        return feat_lengths
+
+    def forward(
+        self,
+        input_values: Optional[torch.Tensor],
+        attention_mask: Optional[torch.Tensor] = None,
+        return_dict: bool = True,
+    ):
+        wavlm_out = self.wavlm(
+            input_values,
+            attention_mask=attention_mask,
+            output_hidden_states=self._use_weighted_layer_sum,
+            return_dict=True,
+        )
+
+        if self._use_weighted_layer_sum:
+            hidden_states = torch.stack(wavlm_out.hidden_states, dim=1)
+            norm_weights = F.softmax(self.layer_weights, dim=0)
+            hidden_states = (hidden_states * norm_weights.view(1, -1, 1, 1)).sum(dim=1)
+        else:
+            hidden_states = wavlm_out.last_hidden_state
+
+        hidden_states = self.projector(hidden_states)
+        for tdnn_layer in self.tdnn:
+            hidden_states = tdnn_layer(hidden_states)
+
+        if attention_mask is None:
+            mean_features = hidden_states.mean(dim=1)
+            std_features = hidden_states.std(dim=1)
+        else:
+            feat_lengths = self.wavlm._get_feat_extract_output_lengths(
+                attention_mask.sum(-1).long()
+            )
+            tdnn_lengths = self._get_tdnn_output_lengths(feat_lengths).clamp(
+                min=1, max=hidden_states.size(1)
+            )
+            T_out = hidden_states.size(1)
+            seq_mask = (
+                torch.arange(T_out, device=hidden_states.device).unsqueeze(0)
+                < tdnn_lengths.unsqueeze(1)
+            ).unsqueeze(-1).float()
+            counts = tdnn_lengths.float().unsqueeze(-1)
+            mean_features = (hidden_states * seq_mask).sum(dim=1) / counts
+            diff = (hidden_states - mean_features.unsqueeze(1)) * seq_mask
+            std_features = (diff.pow(2).sum(dim=1) / counts).sqrt()
+
+        statistic_pooling = torch.cat([mean_features, std_features], dim=-1)
+        output_embeddings = self.feature_extractor(statistic_pooling)
+        logits = self.classifier(output_embeddings)
+
+        return types.SimpleNamespace(embeddings=output_embeddings, logits=logits)
+
+
+class SpeechFaceXVectorEncoder(nn.Module):
+    """
+    Speech encoder for face embedding alignment using WavLMForXVector.
+    """
+
+    def __init__(
+        self,
+        pretrained_path: str = "microsoft/wavlm-base-sv",
+        face_emb_dim: int = 512,
+        dropout: float = 0.1,
+        use_projection: bool = False,
+        projection_hidden_dim: int = None,
+        freeze_feature_encoder: bool = True,
+        reinit_tdnn_projector: bool = False,
+        use_base_wavlm: bool = False,
+        use_weighted_layer_sum: bool = True,
+        shared_projection_config: dict = None,
+        attention_gate_config: dict = None,
+        attribute_disentangled_config: dict = None,
+        attribute_heads_config: dict = None,
+        probabilistic_config: dict = None,
+    ):
+        super().__init__()
+        self.shared_projection_config = shared_projection_config
+
+        if use_base_wavlm:
+            self.wavlm_xvector = WavLMBaseWithXVectorHead(
+                pretrained_path=pretrained_path,
+                use_weighted_layer_sum=use_weighted_layer_sum,
+            )
+        else:
+            logger.info(f"Loading WavLMForXVector from {pretrained_path}")
+            self.wavlm_xvector, loading_info = WavLMForXVector.from_pretrained(
+                pretrained_path, output_loading_info=True
+            )
+            missing = loading_info["missing_keys"]
+            unexpected = loading_info["unexpected_keys"]
+            mismatched = loading_info["mismatched_keys"]
+            if not missing and not unexpected and not mismatched:
+                logger.info("Checkpoint loaded: all keys matched")
+            else:
+                if missing:
+                    logger.warning(f"  Missing keys ({len(missing)}): {missing}")
+                if unexpected:
+                    logger.warning(f"  Unexpected keys ({len(unexpected)}): {unexpected}")
+                if mismatched:
+                    logger.warning(f"  Mismatched shapes ({len(mismatched)}): {mismatched}")
+
+        if reinit_tdnn_projector and not use_base_wavlm:
+            for layer in self.wavlm_xvector.tdnn:
+                layer.kernel.reset_parameters()
+                if hasattr(layer, 'norm'):
+                    layer.norm.reset_parameters()
+            self.wavlm_xvector.projector.reset_parameters()
+
+        self.xvector_dim = self.wavlm_xvector.config.xvector_output_dim
+        self.face_emb_dim = face_emb_dim
+
+        if freeze_feature_encoder:
+            self.wavlm_xvector.freeze_feature_encoder()
+
+        self.use_projection = use_projection or (projection_hidden_dim is not None)
+        self.projection_hidden_dim = projection_hidden_dim
+        if projection_hidden_dim is not None:
+            self.projection = nn.Sequential(
+                nn.Linear(self.xvector_dim, projection_hidden_dim),
+                nn.ReLU(),
+                nn.Linear(projection_hidden_dim, face_emb_dim),
+                nn.BatchNorm1d(face_emb_dim),
+            )
+        elif use_projection and self.xvector_dim != face_emb_dim:
+            self.projection = nn.Sequential(
+                nn.Linear(self.xvector_dim, face_emb_dim),
+                nn.BatchNorm1d(face_emb_dim),
+            )
+        else:
+            self.projection = nn.Identity()
+
+        self.dropout = nn.Dropout(p=dropout) if dropout > 0 else nn.Identity()
+
+        self.attention_gate = None
+        if attention_gate_config and attention_gate_config.get('enabled'):
+            from core.models.encoder.speech_attention_gate import SpeechAttentionGate
+            self.attention_gate = SpeechAttentionGate(
+                embedding_dim=face_emb_dim,
+                hidden_dim=attention_gate_config.get('hidden_dim'),
+                dropout=attention_gate_config.get('dropout', 0.1),
+            )
+
+        self.attribute_disentangled = None
+        self._attr_disent_output = None
+        if attribute_disentangled_config and attribute_disentangled_config.get('enabled'):
+            if self.attention_gate is not None:
+                raise ValueError("attention_gate and attribute_disentangled are mutually exclusive.")
+            from core.models.encoder.attribute_disentangled_encoder import AttributeDisentangledEncoder
+            self.attribute_disentangled = AttributeDisentangledEncoder(
+                embedding_dim=face_emb_dim,
+                num_directions=attribute_disentangled_config.get('num_directions', 16),
+                hidden_dim=attribute_disentangled_config.get('hidden_dim'),
+                dropout=attribute_disentangled_config.get('dropout', 0.1),
+                use_uncertainty=attribute_disentangled_config.get('use_uncertainty', True),
+                use_residual=attribute_disentangled_config.get('use_residual', True),
+            )
+
+        self.attribute_heads = None
+        self._attr_heads_output = None
+        if attribute_heads_config and attribute_heads_config.get('enabled'):
+            from core.models.encoder.attribute_prediction_heads import AttributePredictionHeads
+            self.attribute_heads = AttributePredictionHeads(
+                embedding_dim=face_emb_dim,
+                hidden_dim=attribute_heads_config.get('hidden_dim'),
+                dropout=attribute_heads_config.get('dropout', 0.1),
+            )
+
+        self.probabilistic = False
+        self._probabilistic_output = None
+        if probabilistic_config and probabilistic_config.get('enabled'):
+            self.probabilistic = True
+            self.log_sigma_head = nn.Linear(face_emb_dim, face_emb_dim)
+            self.sigma_min_log = probabilistic_config.get('min_log_sigma', -10.0)
+            self.sigma_max_log = probabilistic_config.get('max_log_sigma', 0.0)
+            nn.init.zeros_(self.log_sigma_head.weight)
+            nn.init.constant_(self.log_sigma_head.bias, self.sigma_min_log)
+
+        self.shared_projection = None
+        if shared_projection_config and shared_projection_config.get('enabled'):
+            from core.models.encoder.projection_head import ProjectionHead
+            self.shared_projection = ProjectionHead(
+                input_dim=face_emb_dim,
+                output_dim=shared_projection_config['shared_dim'],
+                hidden_dim=shared_projection_config.get('hidden_dim'),
+                use_batchnorm=shared_projection_config.get('use_batchnorm', True),
+                dropout=shared_projection_config.get('dropout', 0.1),
+            )
+
+    def freeze_base_model(self):
+        self.wavlm_xvector.freeze_base_model()
+
+    def unfreeze_base_model(self):
+        for param in self.wavlm_xvector.wavlm.parameters():
+            param.requires_grad = True
+
+    def freeze_feature_encoder(self):
+        self.wavlm_xvector.freeze_feature_encoder()
+
+    def get_base_model_params(self) -> List[nn.Parameter]:
+        return [p for p in self.wavlm_xvector.wavlm.parameters() if p.requires_grad]
+
+    def get_head_params(self) -> List[nn.Parameter]:
+        params = []
+        params.extend(self.wavlm_xvector.projector.parameters())
+        for tdnn_layer in self.wavlm_xvector.tdnn:
+            params.extend(tdnn_layer.parameters())
+        params.extend(self.wavlm_xvector.feature_extractor.parameters())
+        if hasattr(self.wavlm_xvector, 'layer_weights'):
+            params.append(self.wavlm_xvector.layer_weights)
+        if self.use_projection:
+            params.extend(self.projection.parameters())
+        if self.attention_gate is not None:
+            params.extend(self.attention_gate.parameters())
+        if self.attribute_disentangled is not None:
+            params.extend(self.attribute_disentangled.parameters())
+        if self.attribute_heads is not None:
+            params.extend(self.attribute_heads.parameters())
+        if self.probabilistic:
+            params.extend(self.log_sigma_head.parameters())
+        if self.shared_projection is not None:
+            params.extend(self.shared_projection.parameters())
+        return [p for p in params if p.requires_grad]
+
+    def get_attr_heads_output(self) -> Optional[Dict[str, torch.Tensor]]:
+        return self._attr_heads_output
+
+    def get_probabilistic_output(self) -> Optional[Dict[str, torch.Tensor]]:
+        return self._probabilistic_output
+
+    def get_kl_loss(self) -> torch.Tensor:
+        if self.attribute_disentangled is None or self._attr_disent_output is None:
+            return torch.tensor(0.0)
+        return self.attribute_disentangled.compute_kl_loss(
+            self._attr_disent_output['alphas'],
+            self._attr_disent_output['log_vars']
+        )
+
+    def get_decorrelation_loss(self) -> torch.Tensor:
+        if self.attribute_disentangled is None or self._attr_disent_output is None:
+            return torch.tensor(0.0)
+        return self.attribute_disentangled.compute_decorrelation_loss(
+            self._attr_disent_output['alphas']
+        )
+
+    def forward(
+        self,
+        audio: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        normalize: bool = True,
+        apply_shared_projection: bool = True,
+    ) -> torch.Tensor:
+        outputs = self.wavlm_xvector(
+            input_values=audio,
+            attention_mask=attention_mask,
+            return_dict=True,
+        )
+
+        embeddings = outputs.embeddings
+        embeddings = self.projection(embeddings)
+        embeddings = self.dropout(embeddings)
+
+        if self.attribute_heads is not None:
+            self._attr_heads_output = self.attribute_heads(embeddings)
+
+        if self.probabilistic:
+            log_sigma = self.log_sigma_head(embeddings)
+            log_sigma = log_sigma.clamp(self.sigma_min_log, self.sigma_max_log)
+            self._probabilistic_output = {'mu': embeddings, 'log_sigma': log_sigma}
+
+        if self.attention_gate is not None:
+            embeddings = self.attention_gate(embeddings)
+            embeddings = F.normalize(embeddings, p=2, dim=1)
+
+        if self.attribute_disentangled is not None:
+            attr_out = self.attribute_disentangled(embeddings)
+            embeddings = attr_out['embedding']
+            self._attr_disent_output = attr_out
+            embeddings = F.normalize(embeddings, p=2, dim=1)
+
+        if self.shared_projection is not None and apply_shared_projection:
+            embeddings = self.shared_projection(embeddings)
+
+        if normalize:
+            embeddings = F.normalize(embeddings, p=2, dim=1)
+
+        return embeddings
+
+    def forward_with_loss(
+        self,
+        audio: torch.Tensor,
+        face_embeddings: torch.Tensor,
+        loss_fn: nn.Module,
+        attention_mask: Optional[torch.Tensor] = None,
+        normalize: bool = True,
+    ) -> Dict[str, torch.Tensor]:
+        speech_embeddings = self.forward(audio, attention_mask, normalize)
+        if normalize:
+            face_embeddings = F.normalize(face_embeddings, p=2, dim=1)
+        loss_dict = loss_fn(speech_embeddings, face_embeddings)
+        loss_dict["speech_embeddings"] = speech_embeddings
+        return loss_dict
+
+    @property
+    def output_dim(self) -> int:
+        if self.shared_projection is not None:
+            return self.shared_projection_config['shared_dim']
+        return self.face_emb_dim
+
+
+def build_speech_face_encoder(
+    pretrained_path: str = "microsoft/wavlm-base-sv",
+    face_emb_dim: int = 512,
+    dropout: float = 0.1,
+    use_projection: bool = True,
+    projection_hidden_dim: int = None,
+    freeze_feature_encoder: bool = True,
+    reinit_tdnn_projector: bool = False,
+    use_base_wavlm: bool = False,
+    use_weighted_layer_sum: bool = True,
+    shared_projection_config: dict = None,
+    attention_gate_config: dict = None,
+    attribute_disentangled_config: dict = None,
+    attribute_heads_config: dict = None,
+    probabilistic_config: dict = None,
+) -> SpeechFaceXVectorEncoder:
+    return SpeechFaceXVectorEncoder(
+        pretrained_path=pretrained_path,
+        face_emb_dim=face_emb_dim,
+        dropout=dropout,
+        use_projection=use_projection,
+        projection_hidden_dim=projection_hidden_dim,
+        freeze_feature_encoder=freeze_feature_encoder,
+        reinit_tdnn_projector=reinit_tdnn_projector,
+        use_base_wavlm=use_base_wavlm,
+        use_weighted_layer_sum=use_weighted_layer_sum,
+        shared_projection_config=shared_projection_config,
+        attention_gate_config=attention_gate_config,
+        attribute_disentangled_config=attribute_disentangled_config,
+        attribute_heads_config=attribute_heads_config,
+        probabilistic_config=probabilistic_config,
+    )

external/arc2face/__init__.py

@@ -0,0 +1,2 @@
+from .models import CLIPTextModelWrapper
+from .utils import project_face_embs, project_face_embs_with_grad, image_align

external/arc2face/models.py

@@ -0,0 +1,83 @@
+import torch
+from transformers import CLIPTextModel
+from typing import Any, Callable, Dict, Optional, Tuple, Union, List
+from transformers.modeling_outputs import BaseModelOutputWithPooling
+from transformers.models.clip.modeling_clip import _make_causal_mask, _expand_mask
+
+
+class CLIPTextModelWrapper(CLIPTextModel):
+    # Adapted from https://github.com/huggingface/transformers/blob/v4.34.1/src/transformers/models/clip/modeling_clip.py#L812
+    # Modified to accept precomputed token embeddings "input_token_embs" as input or calculate them from input_ids and return them.
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        input_token_embs: Optional[torch.Tensor] = None,
+        return_token_embs: Optional[bool] = False,
+    ) -> Union[Tuple, torch.Tensor, BaseModelOutputWithPooling]:
+
+        if return_token_embs:
+            return self.text_model.embeddings.token_embedding(input_ids)
+
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        output_attentions = output_attentions if output_attentions is not None else self.text_model.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.text_model.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.text_model.config.use_return_dict
+
+        if input_ids is None:
+            raise ValueError("You have to specify input_ids")
+
+        input_shape = input_ids.size()
+        input_ids = input_ids.view(-1, input_shape[-1])
+
+        hidden_states = self.text_model.embeddings(input_ids=input_ids, position_ids=position_ids, inputs_embeds=input_token_embs)
+
+        # CLIP's text model uses causal mask, prepare it here.
+        # https://github.com/openai/CLIP/blob/cfcffb90e69f37bf2ff1e988237a0fbe41f33c04/clip/model.py#L324
+        causal_attention_mask = _make_causal_mask(input_shape, hidden_states.dtype, device=hidden_states.device)
+        # expand attention_mask
+        if attention_mask is not None:
+            # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
+            attention_mask = _expand_mask(attention_mask, hidden_states.dtype)
+
+        encoder_outputs = self.text_model.encoder(
+            inputs_embeds=hidden_states,
+            attention_mask=attention_mask,
+            causal_attention_mask=causal_attention_mask,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        last_hidden_state = encoder_outputs[0]
+        last_hidden_state = self.text_model.final_layer_norm(last_hidden_state)
+
+        if self.text_model.eos_token_id == 2:
+            pooled_output = last_hidden_state[
+                torch.arange(last_hidden_state.shape[0], device=last_hidden_state.device),
+                input_ids.to(dtype=torch.int, device=last_hidden_state.device).argmax(dim=-1),
+            ]
+        else:
+            pooled_output = last_hidden_state[
+                torch.arange(last_hidden_state.shape[0], device=last_hidden_state.device),
+                (input_ids.to(dtype=torch.int, device=last_hidden_state.device) == self.text_model.eos_token_id)
+                .int()
+                .argmax(dim=-1),
+            ]
+
+        if not return_dict:
+            return (last_hidden_state, pooled_output) + encoder_outputs[1:]
+
+        return BaseModelOutputWithPooling(
+            last_hidden_state=last_hidden_state,
+            pooler_output=pooled_output,
+            hidden_states=encoder_outputs.hidden_states,
+            attentions=encoder_outputs.attentions,
+        )

external/arc2face/utils.py

@@ -0,0 +1,139 @@
+import scipy
+import PIL
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+@torch.no_grad()
+def project_face_embs(pipeline, face_embs):
+
+    '''
+    face_embs: (N, 512) normalized ArcFace embeddings
+    '''
+
+    arcface_token_id = pipeline.tokenizer.encode("id", add_special_tokens=False)[0]
+
+    input_ids = pipeline.tokenizer(
+            "photo of a id person",
+            truncation=True,
+            padding="max_length",
+            max_length=pipeline.tokenizer.model_max_length,
+            return_tensors="pt",
+        ).input_ids.to(pipeline.device)
+
+    face_embs_padded = F.pad(face_embs, (0, pipeline.text_encoder.config.hidden_size-512), "constant", 0)
+    token_embs = pipeline.text_encoder(input_ids=input_ids.repeat(len(face_embs), 1), return_token_embs=True)
+    token_embs[input_ids==arcface_token_id] = face_embs_padded
+
+    prompt_embeds = pipeline.text_encoder(
+        input_ids=input_ids,
+        input_token_embs=token_embs
+    )[0]
+
+    return prompt_embeds
+
+
+def project_face_embs_with_grad(encoder, tokenizer, face_embs):
+    """
+    Same as project_face_embs but allows gradients for training.
+    """
+    arcface_token_id = tokenizer.encode("id", add_special_tokens=False)[0]
+
+    input_ids = tokenizer(
+        "photo of a id person",
+        truncation=True,
+        padding="max_length",
+        max_length=tokenizer.model_max_length,
+        return_tensors="pt",
+    ).input_ids.to(encoder.device)
+
+    face_embs_padded = F.pad(face_embs, (0, encoder.config.hidden_size - 512), "constant", 0)
+
+    input_ids_batch = input_ids.repeat(len(face_embs), 1)
+    token_embs = encoder(input_ids=input_ids_batch, return_token_embs=True)
+
+    face_embs_padded = face_embs_padded.to(token_embs.dtype)
+    token_embs[input_ids_batch == arcface_token_id] = face_embs_padded
+
+    prompt_embeds = encoder(
+        input_ids=input_ids_batch,
+        input_token_embs=token_embs
+    )[0]
+
+    return prompt_embeds
+
+
+def image_align(img,
+                face_landmarks,
+                output_size=1024,
+                transform_size=4096,
+                enable_padding=True):
+    # Align function from FFHQ dataset pre-processing step
+    # https://github.com/NVlabs/ffhq-dataset/blob/master/download_ffhq.py
+
+    lm = face_landmarks
+    lm_eye_left = lm[36:42]
+    lm_eye_right = lm[42:48]
+    lm_mouth_outer = lm[48:60]
+
+    eye_left = np.mean(lm_eye_left, axis=0)
+    eye_right = np.mean(lm_eye_right, axis=0)
+    eye_avg = (eye_left + eye_right) * 0.5
+    eye_to_eye = eye_right - eye_left
+    mouth_left = lm_mouth_outer[0]
+    mouth_right = lm_mouth_outer[6]
+    mouth_avg = (mouth_left + mouth_right) * 0.5
+    eye_to_mouth = mouth_avg - eye_avg
+
+    x = eye_to_eye - np.flipud(eye_to_mouth) * [-1, 1]
+    x /= np.hypot(*x)
+    x *= max(np.hypot(*eye_to_eye) * 2.0, np.hypot(*eye_to_mouth) * 1.8)
+    y = np.flipud(x) * [-1, 1]
+    c = eye_avg + eye_to_mouth * 0.1
+    quad = np.stack([c - x - y, c - x + y, c + x + y, c + x - y])
+    qsize = np.hypot(*x) * 2
+
+    img = img.convert('RGB')
+
+    shrink = int(np.floor(qsize / output_size * 0.5))
+    if shrink > 1:
+        rsize = (int(np.rint(float(img.size[0]) / shrink)),
+                 int(np.rint(float(img.size[1]) / shrink)))
+        img = img.resize(rsize, PIL.Image.LANCZOS)
+        quad /= shrink
+        qsize /= shrink
+
+    border = max(int(np.rint(qsize * 0.1)), 3)
+    crop = (int(np.floor(min(quad[:, 0]))), int(np.floor(min(quad[:, 1]))),
+            int(np.ceil(max(quad[:, 0]))), int(np.ceil(max(quad[:, 1]))))
+    crop = (max(crop[0] - border, 0), max(crop[1] - border, 0),
+            min(crop[2] + border, img.size[0]), min(crop[3] + border, img.size[1]))
+    if crop[2] - crop[0] < img.size[0] or crop[3] - crop[1] < img.size[1]:
+        img = img.crop(crop)
+        quad -= crop[0:2]
+
+    pad = (int(np.floor(min(quad[:, 0]))), int(np.floor(min(quad[:, 1]))),
+           int(np.ceil(max(quad[:, 0]))), int(np.ceil(max(quad[:, 1]))))
+    pad = (max(-pad[0] + border, 0), max(-pad[1] + border, 0),
+           max(pad[2] - img.size[0] + border, 0), max(pad[3] - img.size[1] + border, 0))
+    if enable_padding and max(pad) > border - 4:
+        pad = np.maximum(pad, int(np.rint(qsize * 0.3)))
+        img = np.pad(np.float32(img),
+                     ((pad[1], pad[3]), (pad[0], pad[2]), (0, 0)), 'reflect')
+        h, w, _ = img.shape
+        y, x, _ = np.ogrid[:h, :w, :1]
+        mask = np.maximum(
+            1.0 - np.minimum(np.float32(x) / pad[0], np.float32(w - 1 - x) / pad[2]),
+            1.0 - np.minimum(np.float32(y) / pad[1], np.float32(h - 1 - y) / pad[3]))
+        blur = qsize * 0.02
+        img += (scipy.ndimage.gaussian_filter(img, [blur, blur, 0]) - img) * np.clip(mask * 3.0 + 1.0, 0.0, 1.0)
+        img += (np.median(img, axis=(0, 1)) - img) * np.clip(mask, 0.0, 1.0)
+        img = PIL.Image.fromarray(np.uint8(np.clip(np.rint(img), 0, 255)), 'RGB')
+        quad += pad[:2]
+
+    img = img.transform((transform_size, transform_size), PIL.Image.QUAD,
+                        (quad + 0.5).flatten(), PIL.Image.BILINEAR)
+    if output_size < transform_size:
+        img = img.resize((output_size, output_size), PIL.Image.LANCZOS)
+
+    return img

requirements.txt

@@ -0,0 +1,17 @@
+gradio
+torch
+torchvision
+torchaudio
+diffusers==0.21.4
+transformers==4.34.1
+accelerate==0.23.0
+peft==0.6.2
+huggingface-hub==0.16.4
+safetensors
+einops
+numpy
+scipy
+pillow
+opencv-python-headless
+facenet-pytorch
+soundfile