Create README.md

README.md

@@ -0,0 +1,139 @@
+---
+tags:
+  - face-restoration
+  - diffusion
+  - one-step
+  - stable-diffusion
+  - lora
+  - image-to-image
+base_model: stabilityai/stable-diffusion-2-1-base
+pipeline_tag: image-to-image
+---
+
+# OSDFace — Pretrained Weights (Mirror)
+
+
+![github.com_jkwang28_OSDFace_](https://cdn-uploads.huggingface.co/production/uploads/6557b9b4deee83130ac92941/Qs7aOIkhOyghdH8Bdt5wF.png)
+
+> **This is an unofficial mirror.**
+> All credit goes to the original authors. The weights are mirrored here from the [official OSDFace repository](https://github.com/jkwang28/OSDFace) for convenience, as the original download is hosted on OneDrive/Google Drive which can be slow or inaccessible in some regions.
+> Please cite the original paper and star the original repo if you use these weights.
+
+## Overview
+
+OSDFace (**One-Step Diffusion Model for Face Restoration**) is a single-step diffusion model that restores degraded, low-quality face images into high-fidelity, identity-consistent outputs. It was accepted at **CVPR 2025**.
+
+Unlike multi-step diffusion approaches, OSDFace requires only **one forward pass** through a modified Stable Diffusion 2.1 UNet, making it significantly faster at inference while achieving state-of-the-art results on both synthetic (CelebA-Test) and real-world (Wider-Test, LFW-Test, WebPhoto-Test) benchmarks.
+
+The key innovations are:
+
+- **Visual Representation Embedder (VRE):** A VQ-VAE encoder that tokenizes the low-quality input face and produces visual prompt embeddings via a vector-quantized dictionary. These embeddings replace the text encoder's output and are fed directly into the UNet's cross-attention layers.
+- **Facial Identity Loss:** A face-recognition-derived loss that enforces identity consistency between the restored and ground-truth faces.
+- **GAN Guidance:** A generative adversarial network guides the one-step diffusion to align the output distribution with the ground truth.
+
+## Usage
+
+### Prerequisites
+
+- **Base model:** [stabilityai/stable-diffusion-2-1-base](https://huggingface.co/stabilityai/stable-diffusion-2-1-base)
+- **Python 3.10**, PyTorch 2.4.0, diffusers 0.27.2
+
+### Quick Start
+
+```bash
+# Clone the official repo
+git clone https://github.com/jkwang28/OSDFace.git
+cd OSDFace
+
+# Download these weights into pretrained/
+# Place: associate_2.ckpt, embedding_change_weights.pth, pytorch_lora_weights.safetensors
+
+# Run inference (with LoRA merging for speed)
+python infer.py \
+    --input_image data/WebPhoto-Test \
+    --output_dir results/WebPhoto-Test \
+    --pretrained_model_name_or_path "stabilityai/stable-diffusion-2-1-base" \
+    --img_encoder_weight "pretrained/associate_2.ckpt" \
+    --ckpt_path pretrained \
+    --merge_lora \
+    --mixed_precision fp16 \
+    --gpu_ids 0
+```
+
+> **Note on the different pretrained model**
+> Although the project is based on `stabilityai/stable-diffusion-2-1-base` we use `Manojb/stable-diffusion-2-1-base` because the former can't be downloaded from huggingface.
+
+
+## Files in This Repository
+
+![image](https://cdn-uploads.huggingface.co/production/uploads/6557b9b4deee83130ac92941/rqOY2DHGJ1MhnJkJ5keJ0.png)
+
+### `associate_2.ckpt` (1.87 GB)
+
+The **VQ-VAE image encoder** (referred to as the Visual Representation Embedder in the paper). This is the core component that understands the degraded input face.
+
+It contains a multi-head encoder with downsampling blocks, a mid-block with attention, and a vector quantizer with a learned 1024-entry codebook (embedding dim 512). At inference, the encoder processes a 512×512 low-quality face, extracts spatial features, quantizes them against the codebook, and selects the 77 closest (non-duplicate) codebook entries — producing a `(batch, 77, 512)` tensor that acts as a drop-in replacement for CLIP text embeddings in the UNet's cross-attention.
+
+**Loaded via:** `--img_encoder_weight associate_2.ckpt`
+
+### `embedding_change_weights.pth` (1.58 MB)
+
+A lightweight **embedding projection module** (`TwoLayerConv1x1`) that maps the VRE output from 512 dimensions to 1024 dimensions, matching the hidden size expected by Stable Diffusion 2.1's UNet cross-attention layers.
+
+Architecture: two 1×1 Conv1d layers with SiLU activations (`512 → 256 → 1024`), operating over the 77-token sequence.
+
+This module is used in the default configuration (without `--cat_prompt_embedding`). When `--cat_prompt_embedding` is enabled, the VRE instead outputs 154 tokens at 512-dim which are reshaped to 77 tokens at 1024-dim, bypassing this module entirely.
+
+**Loaded from:** `<ckpt_path>/embedding_change_weights.pth`
+
+### `pytorch_lora_weights.safetensors` (67.9 MB)
+
+**LoRA (Low-Rank Adaptation) weights** for the Stable Diffusion 2.1 UNet. These adapt the frozen SD2.1 UNet to perform one-step face restoration conditioned on the VRE embeddings.
+
+Default LoRA configuration: **rank 16, alpha 16** (effective scaling factor `alpha/rank = 1.0`). The weights cover both standard LoRA layers (`lora_A`/`lora_B`) and some additional `lora.up`/`lora.down` layers.
+
+These can be loaded in two ways:
+- **Dynamic loading** (default): loaded at runtime via `diffusers`' `load_lora_weights()`
+- **Merged loading** (`--merge_lora`): pre-merged into the UNet weights before inference for slightly faster execution
+
+**Loaded from:** `<ckpt_path>/pytorch_lora_weights.safetensors`
+
+
+## Key Inference Arguments
+
+| Argument | Default | Description |
+|---|---|---|
+| `--merge_lora` | off | Merge LoRA into UNet weights (recommended) |
+| `--mixed_precision` | `fp32` | Use `fp16` for faster inference / lower VRAM |
+| `--gpu_ids` | `[0]` | Multi-GPU support, e.g. `--gpu_ids 0 1 2 3` |
+| `--cat_prompt_embedding` | off | Alternative embedding strategy (skips embedding_change module) |
+| `--lora_rank` | 16 | LoRA rank (must match training) |
+| `--lora_alpha` | 16 | LoRA alpha (must match training) |
+
+## Inference Pipeline (Summary)
+
+1. Input image resized to **512×512**
+2. VRE encodes the LQ face → `(B, 77, 512)` visual prompt
+3. Embedding projection maps to `(B, 77, 1024)` (or concatenation path)
+4. VAE encodes the LQ face to latent space
+5. UNet performs a **single denoising step** at timestep 399, conditioned on the visual prompt
+6. Predicted clean latent is decoded by the VAE → restored face
+
+## Citation
+
+```bibtex
+@InProceedings{wang2025osdface,
+    author    = {Wang, Jingkai and Gong, Jue and Zhang, Lin and Chen, Zheng and Liu, Xing and Gu, Hong and Liu, Yutong and Zhang, Yulun and Yang, Xiaokang},
+    title     = {{OSDFace}: One-Step Diffusion Model for Face Restoration},
+    booktitle = {Proceedings of the Computer Vision and Pattern Recognition Conference (CVPR)},
+    month     = {June},
+    year      = {2025},
+    pages     = {12626-12636}
+}
+```
+
+## Links
+
+- 📄 [Paper (arXiv)](https://arxiv.org/abs/2411.17163)
+- 💻 [Official Repository](https://github.com/jkwang28/OSDFace)
+- 🌐 [Project Page](https://www.jingkaiwang.com/OSDFace/)