new

.agents/skills/uv-package-manager/SKILL.md

@@ -1,834 +0,0 @@
----
-name: uv-package-manager
-description: Master the uv package manager for fast Python dependency management, virtual environments, and modern Python project workflows. Use when setting up Python projects, managing dependencies, or optimizing Python development workflows with uv.
----
-
-# UV Package Manager
-
-Comprehensive guide to using uv, an extremely fast Python package installer and resolver written in Rust, for modern Python project management and dependency workflows.
-
-## When to Use This Skill
-
-- Setting up new Python projects quickly
-- Managing Python dependencies faster than pip
-- Creating and managing virtual environments
-- Installing Python interpreters
-- Resolving dependency conflicts efficiently
-- Migrating from pip/pip-tools/poetry
-- Speeding up CI/CD pipelines
-- Managing monorepo Python projects
-- Working with lockfiles for reproducible builds
-- Optimizing Docker builds with Python dependencies
-
-## Core Concepts
-
-### 1. What is uv?
-
-- **Ultra-fast package installer**: 10-100x faster than pip
-- **Written in Rust**: Leverages Rust's performance
-- **Drop-in pip replacement**: Compatible with pip workflows
-- **Virtual environment manager**: Create and manage venvs
-- **Python installer**: Download and manage Python versions
-- **Resolver**: Advanced dependency resolution
-- **Lockfile support**: Reproducible installations
-
-### 2. Key Features
-
-- Blazing fast installation speeds
-- Disk space efficient with global cache
-- Compatible with pip, pip-tools, poetry
-- Comprehensive dependency resolution
-- Cross-platform support (Linux, macOS, Windows)
-- No Python required for installation
-- Built-in virtual environment support
-
-### 3. UV vs Traditional Tools
-
-- **vs pip**: 10-100x faster, better resolver
-- **vs pip-tools**: Faster, simpler, better UX
-- **vs poetry**: Faster, less opinionated, lighter
-- **vs conda**: Faster, Python-focused
-
-## Installation
-
-### Quick Install
-
-```bash
-# macOS/Linux
-curl -LsSf https://astral.sh/uv/install.sh | sh
-
-# Windows (PowerShell)
-powershell -c "irm https://astral.sh/uv/install.ps1 | iex"
-
-# Using pip (if you already have Python)
-pip install uv
-
-# Using Homebrew (macOS)
-brew install uv
-
-# Using cargo (if you have Rust)
-cargo install --git https://github.com/astral-sh/uv uv
-```
-
-### Verify Installation
-
-```bash
-uv --version
-# uv 0.x.x
-```
-
-## Quick Start
-
-### Create a New Project
-
-```bash
-# Create new project with virtual environment
-uv init my-project
-cd my-project
-
-# Or create in current directory
-uv init .
-
-# Initialize creates:
-# - .python-version (Python version)
-# - pyproject.toml (project config)
-# - README.md
-# - .gitignore
-```
-
-### Install Dependencies
-
-```bash
-# Install packages (creates venv if needed)
-uv add requests pandas
-
-# Install dev dependencies
-uv add --dev pytest black ruff
-
-# Install from requirements.txt
-uv pip install -r requirements.txt
-
-# Install from pyproject.toml
-uv sync
-```
-
-## Virtual Environment Management
-
-### Pattern 1: Creating Virtual Environments
-
-```bash
-# Create virtual environment with uv
-uv venv
-
-# Create with specific Python version
-uv venv --python 3.12
-
-# Create with custom name
-uv venv my-env
-
-# Create with system site packages
-uv venv --system-site-packages
-
-# Specify location
-uv venv /path/to/venv
-```
-
-### Pattern 2: Activating Virtual Environments
-
-```bash
-# Linux/macOS
-source .venv/bin/activate
-
-# Windows (Command Prompt)
-.venv\Scripts\activate.bat
-
-# Windows (PowerShell)
-.venv\Scripts\Activate.ps1
-
-# Or use uv run (no activation needed)
-uv run python script.py
-uv run pytest
-```
-
-### Pattern 3: Using uv run
-
-```bash
-# Run Python script (auto-activates venv)
-uv run python app.py
-
-# Run installed CLI tool
-uv run black .
-uv run pytest
-
-# Run with specific Python version
-uv run --python 3.11 python script.py
-
-# Pass arguments
-uv run python script.py --arg value
-```
-
-## Package Management
-
-### Pattern 4: Adding Dependencies
-
-```bash
-# Add package (adds to pyproject.toml)
-uv add requests
-
-# Add with version constraint
-uv add "django>=4.0,<5.0"
-
-# Add multiple packages
-uv add numpy pandas matplotlib
-
-# Add dev dependency
-uv add --dev pytest pytest-cov
-
-# Add optional dependency group
-uv add --optional docs sphinx
-
-# Add from git
-uv add git+https://github.com/user/repo.git
-
-# Add from git with specific ref
-uv add git+https://github.com/user/repo.git@v1.0.0
-
-# Add from local path
-uv add ./local-package
-
-# Add editable local package
-uv add -e ./local-package
-```
-
-### Pattern 5: Removing Dependencies
-
-```bash
-# Remove package
-uv remove requests
-
-# Remove dev dependency
-uv remove --dev pytest
-
-# Remove multiple packages
-uv remove numpy pandas matplotlib
-```
-
-### Pattern 6: Upgrading Dependencies
-
-```bash
-# Upgrade specific package
-uv add --upgrade requests
-
-# Upgrade all packages
-uv sync --upgrade
-
-# Upgrade package to latest
-uv add --upgrade requests
-
-# Show what would be upgraded
-uv tree --outdated
-```
-
-### Pattern 7: Locking Dependencies
-
-```bash
-# Generate uv.lock file
-uv lock
-
-# Update lock file
-uv lock --upgrade
-
-# Lock without installing
-uv lock --no-install
-
-# Lock specific package
-uv lock --upgrade-package requests
-```
-
-## Python Version Management
-
-### Pattern 8: Installing Python Versions
-
-```bash
-# Install Python version
-uv python install 3.12
-
-# Install multiple versions
-uv python install 3.11 3.12 3.13
-
-# Install latest version
-uv python install
-
-# List installed versions
-uv python list
-
-# Find available versions
-uv python list --all-versions
-```
-
-### Pattern 9: Setting Python Version
-
-```bash
-# Set Python version for project
-uv python pin 3.12
-
-# This creates/updates .python-version file
-
-# Use specific Python version for command
-uv --python 3.11 run python script.py
-
-# Create venv with specific version
-uv venv --python 3.12
-```
-
-## Project Configuration
-
-### Pattern 10: pyproject.toml with uv
-
-```toml
-[project]
-name = "my-project"
-version = "0.1.0"
-description = "My awesome project"
-readme = "README.md"
-requires-python = ">=3.8"
-dependencies = [
-    "requests>=2.31.0",
-    "pydantic>=2.0.0",
-    "click>=8.1.0",
-]
-
-[project.optional-dependencies]
-dev = [
-    "pytest>=7.4.0",
-    "pytest-cov>=4.1.0",
-    "black>=23.0.0",
-    "ruff>=0.1.0",
-    "mypy>=1.5.0",
-]
-docs = [
-    "sphinx>=7.0.0",
-    "sphinx-rtd-theme>=1.3.0",
-]
-
-[build-system]
-requires = ["hatchling"]
-build-backend = "hatchling.build"
-
-[tool.uv]
-dev-dependencies = [
-    # Additional dev dependencies managed by uv
-]
-
-[tool.uv.sources]
-# Custom package sources
-my-package = { git = "https://github.com/user/repo.git" }
-```
-
-### Pattern 11: Using uv with Existing Projects
-
-```bash
-# Migrate from requirements.txt
-uv add -r requirements.txt
-
-# Migrate from poetry
-# Already have pyproject.toml, just use:
-uv sync
-
-# Export to requirements.txt
-uv pip freeze > requirements.txt
-
-# Export with hashes
-uv pip freeze --require-hashes > requirements.txt
-```
-
-## Advanced Workflows
-
-### Pattern 12: Monorepo Support
-
-```bash
-# Project structure
-# monorepo/
-#   packages/
-#     package-a/
-#       pyproject.toml
-#     package-b/
-#       pyproject.toml
-#   pyproject.toml (root)
-
-# Root pyproject.toml
-[tool.uv.workspace]
-members = ["packages/*"]
-
-# Install all workspace packages
-uv sync
-
-# Add workspace dependency
-uv add --path ./packages/package-a
-```
-
-### Pattern 13: CI/CD Integration
-
-```yaml
-# .github/workflows/test.yml
-name: Tests
-
-on: [push, pull_request]
-
-jobs:
-  test:
-    runs-on: ubuntu-latest
-
-    steps:
-      - uses: actions/checkout@v4
-
-      - name: Install uv
-        uses: astral-sh/setup-uv@v2
-        with:
-          enable-cache: true
-
-      - name: Set up Python
-        run: uv python install 3.12
-
-      - name: Install dependencies
-        run: uv sync --all-extras --dev
-
-      - name: Run tests
-        run: uv run pytest
-
-      - name: Run linting
-        run: |
-          uv run ruff check .
-          uv run black --check .
-```
-
-### Pattern 14: Docker Integration
-
-```dockerfile
-# Dockerfile
-FROM python:3.12-slim
-
-# Install uv
-COPY --from=ghcr.io/astral-sh/uv:latest /uv /usr/local/bin/uv
-
-# Set working directory
-WORKDIR /app
-
-# Copy dependency files
-COPY pyproject.toml uv.lock ./
-
-# Install dependencies
-RUN uv sync --frozen --no-dev
-
-# Copy application code
-COPY . .
-
-# Run application
-CMD ["uv", "run", "python", "app.py"]
-```
-
-**Optimized multi-stage build:**
-
-```dockerfile
-# Multi-stage Dockerfile
-FROM python:3.12-slim AS builder
-
-# Install uv
-COPY --from=ghcr.io/astral-sh/uv:latest /uv /usr/local/bin/uv
-
-WORKDIR /app
-
-# Install dependencies to venv
-COPY pyproject.toml uv.lock ./
-RUN uv sync --frozen --no-dev --no-editable
-
-# Runtime stage
-FROM python:3.12-slim
-
-WORKDIR /app
-
-# Copy venv from builder
-COPY --from=builder /app/.venv .venv
-COPY . .
-
-# Use venv
-ENV PATH="/app/.venv/bin:$PATH"
-
-CMD ["python", "app.py"]
-```
-
-### Pattern 15: Lockfile Workflows
-
-```bash
-# Create lockfile (uv.lock)
-uv lock
-
-# Install from lockfile (exact versions)
-uv sync --frozen
-
-# Update lockfile without installing
-uv lock --no-install
-
-# Upgrade specific package in lock
-uv lock --upgrade-package requests
-
-# Check if lockfile is up to date
-uv lock --check
-
-# Export lockfile to requirements.txt
-uv export --format requirements-txt > requirements.txt
-
-# Export with hashes for security
-uv export --format requirements-txt --hash > requirements.txt
-```
-
-## Performance Optimization
-
-### Pattern 16: Using Global Cache
-
-```bash
-# UV automatically uses global cache at:
-# Linux: ~/.cache/uv
-# macOS: ~/Library/Caches/uv
-# Windows: %LOCALAPPDATA%\uv\cache
-
-# Clear cache
-uv cache clean
-
-# Check cache size
-uv cache dir
-```
-
-### Pattern 17: Parallel Installation
-
-```bash
-# UV installs packages in parallel by default
-
-# Control parallelism
-uv pip install --jobs 4 package1 package2
-
-# No parallel (sequential)
-uv pip install --jobs 1 package
-```
-
-### Pattern 18: Offline Mode
-
-```bash
-# Install from cache only (no network)
-uv pip install --offline package
-
-# Sync from lockfile offline
-uv sync --frozen --offline
-```
-
-## Comparison with Other Tools
-
-### uv vs pip
-
-```bash
-# pip
-python -m venv .venv
-source .venv/bin/activate
-pip install requests pandas numpy
-# ~30 seconds
-
-# uv
-uv venv
-uv add requests pandas numpy
-# ~2 seconds (10-15x faster)
-```
-
-### uv vs poetry
-
-```bash
-# poetry
-poetry init
-poetry add requests pandas
-poetry install
-# ~20 seconds
-
-# uv
-uv init
-uv add requests pandas
-uv sync
-# ~3 seconds (6-7x faster)
-```
-
-### uv vs pip-tools
-
-```bash
-# pip-tools
-pip-compile requirements.in
-pip-sync requirements.txt
-# ~15 seconds
-
-# uv
-uv lock
-uv sync --frozen
-# ~2 seconds (7-8x faster)
-```
-
-## Common Workflows
-
-### Pattern 19: Starting a New Project
-
-```bash
-# Complete workflow
-uv init my-project
-cd my-project
-
-# Set Python version
-uv python pin 3.12
-
-# Add dependencies
-uv add fastapi uvicorn pydantic
-
-# Add dev dependencies
-uv add --dev pytest black ruff mypy
-
-# Create structure
-mkdir -p src/my_project tests
-
-# Run tests
-uv run pytest
-
-# Format code
-uv run black .
-uv run ruff check .
-```
-
-### Pattern 20: Maintaining Existing Project
-
-```bash
-# Clone repository
-git clone https://github.com/user/project.git
-cd project
-
-# Install dependencies (creates venv automatically)
-uv sync
-
-# Install with dev dependencies
-uv sync --all-extras
-
-# Update dependencies
-uv lock --upgrade
-
-# Run application
-uv run python app.py
-
-# Run tests
-uv run pytest
-
-# Add new dependency
-uv add new-package
-
-# Commit updated files
-git add pyproject.toml uv.lock
-git commit -m "Add new-package dependency"
-```
-
-## Tool Integration
-
-### Pattern 21: Pre-commit Hooks
-
-```yaml
-# .pre-commit-config.yaml
-repos:
-  - repo: local
-    hooks:
-      - id: uv-lock
-        name: uv lock
-        entry: uv lock
-        language: system
-        pass_filenames: false
-
-      - id: ruff
-        name: ruff
-        entry: uv run ruff check --fix
-        language: system
-        types: [python]
-
-      - id: black
-        name: black
-        entry: uv run black
-        language: system
-        types: [python]
-```
-
-### Pattern 22: VS Code Integration
-
-```json
-// .vscode/settings.json
-{
-  "python.defaultInterpreterPath": "${workspaceFolder}/.venv/bin/python",
-  "python.terminal.activateEnvironment": true,
-  "python.testing.pytestEnabled": true,
-  "python.testing.pytestArgs": ["-v"],
-  "python.linting.enabled": true,
-  "python.formatting.provider": "black",
-  "[python]": {
-    "editor.defaultFormatter": "ms-python.black-formatter",
-    "editor.formatOnSave": true
-  }
-}
-```
-
-## Troubleshooting
-
-### Common Issues
-
-```bash
-# Issue: uv not found
-# Solution: Add to PATH or reinstall
-echo 'export PATH="$HOME/.cargo/bin:$PATH"' >> ~/.bashrc
-
-# Issue: Wrong Python version
-# Solution: Pin version explicitly
-uv python pin 3.12
-uv venv --python 3.12
-
-# Issue: Dependency conflict
-# Solution: Check resolution
-uv lock --verbose
-
-# Issue: Cache issues
-# Solution: Clear cache
-uv cache clean
-
-# Issue: Lockfile out of sync
-# Solution: Regenerate
-uv lock --upgrade
-```
-
-## Best Practices
-
-### Project Setup
-
-1. **Always use lockfiles** for reproducibility
-2. **Pin Python version** with .python-version
-3. **Separate dev dependencies** from production
-4. **Use uv run** instead of activating venv
-5. **Commit uv.lock** to version control
-6. **Use --frozen in CI** for consistent builds
-7. **Leverage global cache** for speed
-8. **Use workspace** for monorepos
-9. **Export requirements.txt** for compatibility
-10. **Keep uv updated** for latest features
-
-### Performance Tips
-
-```bash
-# Use frozen installs in CI
-uv sync --frozen
-
-# Use offline mode when possible
-uv sync --offline
-
-# Parallel operations (automatic)
-# uv does this by default
-
-# Reuse cache across environments
-# uv shares cache globally
-
-# Use lockfiles to skip resolution
-uv sync --frozen  # skips resolution
-```
-
-## Migration Guide
-
-### From pip + requirements.txt
-
-```bash
-# Before
-python -m venv .venv
-source .venv/bin/activate
-pip install -r requirements.txt
-
-# After
-uv venv
-uv pip install -r requirements.txt
-# Or better:
-uv init
-uv add -r requirements.txt
-```
-
-### From Poetry
-
-```bash
-# Before
-poetry install
-poetry add requests
-
-# After
-uv sync
-uv add requests
-
-# Keep existing pyproject.toml
-# uv reads [project] and [tool.poetry] sections
-```
-
-### From pip-tools
-
-```bash
-# Before
-pip-compile requirements.in
-pip-sync requirements.txt
-
-# After
-uv lock
-uv sync --frozen
-```
-
-## Command Reference
-
-### Essential Commands
-
-```bash
-# Project management
-uv init [PATH]              # Initialize project
-uv add PACKAGE              # Add dependency
-uv remove PACKAGE           # Remove dependency
-uv sync                     # Install dependencies
-uv lock                     # Create/update lockfile
-
-# Virtual environments
-uv venv [PATH]              # Create venv
-uv run COMMAND              # Run in venv
-
-# Python management
-uv python install VERSION   # Install Python
-uv python list              # List installed Pythons
-uv python pin VERSION       # Pin Python version
-
-# Package installation (pip-compatible)
-uv pip install PACKAGE      # Install package
-uv pip uninstall PACKAGE    # Uninstall package
-uv pip freeze               # List installed
-uv pip list                 # List packages
-
-# Utility
-uv cache clean              # Clear cache
-uv cache dir                # Show cache location
-uv --version                # Show version
-```
-
-## Resources
-
-- **Official documentation**: https://docs.astral.sh/uv/
-- **GitHub repository**: https://github.com/astral-sh/uv
-- **Astral blog**: https://astral.sh/blog
-- **Migration guides**: https://docs.astral.sh/uv/guides/
-- **Comparison with other tools**: https://docs.astral.sh/uv/pip/compatibility/
-
-## Best Practices Summary
-
-1. **Use uv for all new projects** - Start with `uv init`
-2. **Commit lockfiles** - Ensure reproducible builds
-3. **Pin Python versions** - Use .python-version
-4. **Use uv run** - Avoid manual venv activation
-5. **Leverage caching** - Let uv manage global cache
-6. **Use --frozen in CI** - Exact reproduction
-7. **Keep uv updated** - Fast-moving project
-8. **Use workspaces** - For monorepo projects
-9. **Export for compatibility** - Generate requirements.txt when needed
-10. **Read the docs** - uv is feature-rich and evolving

.claude/skills/uv-package-manager

@@ -1 +0,0 @@
-../../.agents/skills/uv-package-manager

.dockerignore

@@ -1,38 +0,0 @@
-# Python
-__pycache__/
-*.py[cod]
-*$py.class
-*.so
-.Python
-*.egg-info/
-dist/
-build/
-
-# Virtual Environment
-.venv/
-venv/
-
-# Cache
-.cache/
-*.log
-*.tmp
-
-# Output directories
-processed_results/
-
-# IDE
-.vscode/
-.idea/
-*.swp
-
-# Git
-.git/
-.gitignore
-
-# Documentation
-*.md
-!README.md
-
-# Test
-tests/
-eval/

.github/workflows/docker.yml

@@ -1,57 +0,0 @@
-name: Build and Push to GHCR
-
-on:
-  push:
-    branches:
-      - main
-    tags:
-      - 'v*'
-  workflow_dispatch:
-
-env:
-  REGISTRY: ghcr.io
-  IMAGE_NAME: ${{ github.repository }}
-
-jobs:
-  build:
-    runs-on: self-hosted
-    permissions:
-      contents: read
-      packages: write
-
-    steps:
-      - name: Checkout
-        uses: actions/checkout@v4
-
-      - name: Set up Docker Buildx
-        uses: docker/setup-buildx-action@v3
-
-      - name: Log in to GHCR
-        uses: docker/login-action@v3
-        with:
-          registry: ${{ env.REGISTRY }}
-          username: ${{ github.actor }}
-          password: ${{ secrets.GITHUB_TOKEN }}
-
-      - name: Extract metadata
-        id: meta
-        uses: docker/metadata-action@v5
-        with:
-          images: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}
-          tags: |
-            type=ref,event=branch
-            type=ref,event=pr
-            type=semver,pattern={{version}}
-            type=semver,pattern={{major}}.{{minor}}
-            type=sha,prefix=
-
-      - name: Build and push
-        uses: docker/build-push-action@v5
-        with:
-          context: .
-          push: true
-          tags: ${{ steps.meta.outputs.tags }}
-          labels: ${{ steps.meta.outputs.labels }}
-          platforms: linux/amd64
-          cache-from: type=gha
-          cache-to: type=gha,mode=max

.gitignore

@@ -1,21 +0,0 @@
-checkpoints/**/*.pth
-checkpoints/**/*.pt
-checkpoints/**/*.pkl
-checkpoints/**/*.zip
-checkpoints/**/*.safetensors
-checkpoints/**/*
-
-# Ignore local dependencies (installed via pip/uv)
-# latentsync/ - Keep for HuggingFace Spaces
-# tigersound/
-# FastAudioSR/
-# descript-audiotools/
-# models/
-
-# Python cache and virtual environment
-__pycache__/
-*.pyc
-.pytest_cache/
-.coverage
-.venv/
-uv.lock

.opencode/skills/uv-package-manager

@@ -1 +0,0 @@
-../../.agents/skills/uv-package-manager

.python-version

@@ -1 +0,0 @@
-3.10

AGENTS.md

@@ -1,226 +0,0 @@
-# AGENTS.md
-
-This file provides guidance for agentic coding assistants working in this repository.
-
-## Build/Install Commands
-
-### HuggingFace Spaces (Deployment)
-```bash
-# HuggingFace automatically runs:
-pip install -r requirements.txt
-python app.py
-```
-
-### Local Development (with uv)
-```bash
-# Setup (one-time)
-./setup_local.sh
-
-# Or manually:
-uv venv
-uv pip install -r requirements.txt
-
-# Run app
-uv run python app.py
-```
-
-### Local Development (with pip - standard)
-```bash
-python3 -m venv .venv
-source .venv/bin/activate
-pip install -r requirements.txt
-python app.py
-```
-
-**Note**: Use uv for faster dependency installation (10-100x faster than pip). This project currently has NO linting, formatting, or test infrastructure configured. When making changes, ensure the application runs successfully with `uv run python app.py` (or `python app.py` after activating venv).
-
----
-
-## Application Architecture
-
-This project provides lipsync functionality for English videos:
-
-- **Lipsync Only** (`lipsync_only_video`): Apply lipsync to English video
-
-### Core Workflow:
-
-```
-Upload Video → Crop Duration → Extract Audio (vocal/bg/effect) → Upsample Audio →
-Lipsync → Merge Video + Audio
-```
-
----
-
-## Code Style Guidelines
-
-### Imports
-
-- Organize imports in this order:
-  1. Standard library (os, sys, pathlib, subprocess, etc.)
-  2. Third-party packages (torch, gradio, numpy, etc.)
-  3. Local/custom modules (lipsync, time_util)
-- Use absolute imports for clarity
-- Keep all imports at the top of files (avoid scattered imports)
-- Remove duplicate imports
-
-**Example:**
-```python
-import os
-import subprocess
-from pathlib import Path
-from typing import List, Dict
-
-import torch
-import gradio as gr
-from pydub import AudioSegment
-
-from lipsync import apply_lipsync
-from time_util import timer
-```
-
----
-
-### Formatting & Spacing
-
-- Use 4 spaces for indentation (not tabs)
-- Use f-strings for string formatting
-- Keep lines under 100 characters where practical
-- Add 2 blank lines before top-level function definitions
-
----
-
-### Type Hints
-
-- Use type hints for function signatures when clear
-- Use `| None` for optional types (Python 3.10+) instead of `Optional[T]`
-
-**Example:**
-```python
-def format_timestamp(ts: float) -> str:
-    """Convert seconds to SRT timestamp format."""
-    hrs = int(ts // 3600)
-    # ...
-```
-
----
-
-### Naming Conventions
-
-- **Functions & variables**: `snake_case`
-- **Constants**: `UPPER_CASE`
-- **Classes**: `PascalCase`
-
-**Example:**
-```python
-MODEL_SIZE = "medium"
-MAX_BATCH_MS = 300_000
-
-def extract_audio_to_wav(input_video: str, output_dir: str):
-    pass
-```
-
----
-
-### Error Handling
-
-- Use `subprocess.check_call()` or `subprocess.run(check=True)` for shell commands
-- Use `gr.Error()` for user-facing errors in Gradio callbacks
-- Include descriptive error messages
-
-**Example:**
-```python
-def lipsync_only_video(video_file, duration, session_id=None, progress=None):
-    if video_file is None:
-        raise gr.Error("Please upload a clip.")
-```
-
----
-
-### Docstrings
-
-- Add brief docstrings for non-trivial functions explaining purpose and parameters
-- Keep docstrings concise
-
----
-
-### Context Managers
-
-- Use `@contextmanager` decorators for resource management when appropriate
-- Use `with` statements for file operations and subprocess calls
-
-**Example:**
-```python
-from contextlib import contextmanager
-
-@contextmanager
-def timer(name: str):
-    start = time.time()
-    print(f"{name}...")
-    yield
-    print(f"  -> {name} completed in {time.time() - start:.2f} sec")
-```
-
----
-
-### Logging
-
-- Use the `logging` module over print() for production code
-- Configure log levels appropriately
-
-**Example:**
-```python
-import logging
-
-logging.getLogger("httpx").setLevel(logging.WARNING)
-logging.getLogger("httpcore").setLevel(logging.WARNING)
-```
-
----
-
-### File Paths
-
-- Use `pathlib.Path` for cross-platform path handling when possible
-- Use `os.path.join()` for constructing paths
-- Use `os.makedirs(path, exist_ok=True)` for directory creation
-
----
-
-### PyTorch Best Practices
-
-- Use `torch.no_grad()` context manager for inference
-- Move tensors to appropriate device explicitly: `.to("cuda")`
-- Clear GPU cache after operations: `torch.cuda.empty_cache()`
-- Set seed for reproducibility when needed: `torch.manual_seed(1234)`
-
----
-
-### Gradio Integration
-
-- Use `gr.Progress(track_tqdm=True)` for progress tracking
-- Return appropriate outputs: tuples matching the defined outputs
-- Handle session state carefully for multi-user environments
-
----
-
-## Key Helper Functions
-
-The application uses modular helper functions for better maintainability:
-
-- **`setup_output_dir(session_id)`**: Creates output directory for session
-- **`crop_video_duration(video_path, duration, output_dir)`**: Crops video using FFmpeg
-- **`extract_audio_to_wav(video_path, output_dir)`**: Extracts and separates audio tracks
-- **`upsample_audio(audio_path, output_dir)`**: Upsamples 16kHz audio to 48kHz
-- **`merge_audio_video(video_path, audio_path, output_dir)`**: Merges video with audio
-- **`apply_lipsync_to_video(video_path, audio_16k_path, output_dir)`**: Wraps lipsync.apply_lipsync
-
----
-
-## Important Notes
-
-- **CUDA Required**: This project requires GPU (CUDA) for most ML operations
-- **Large Models**: Models are loaded at startup; avoid unnecessary reloading
-- **Session Management**: Each user session creates a unique output directory that should be cleaned up
-- **Clean Code**: Each function should have a single responsibility - prefer creating new helper functions over adding complexity to existing ones
-- **Flash-attn-3**: Only available for Linux x86_64. Comment out in requirements.txt for local testing on macOS
-- **Dependencies Managed via uv (local) / pip (HuggingFace)**: All packages are installed in `.venv/` (gitignored)
-- **Local Dependencies Ignored**: latentsync/, tigersound/, FastAudioSR/, descript-audiotools/ directories are gitignored - packages are installed from git repos via requirements.txt

Dockerfile

@@ -1,48 +0,0 @@
-FROM docker.io/nvidia/cuda:12.3.2-cudnn9-devel-ubuntu22.04 AS builder
-
-ENV DEBIAN_FRONTEND=noninteractive
-
-RUN apt-get update && apt-get install -y --no-install-recommends \
-	git rsync \
-	make build-essential libssl-dev zlib1g-dev \
-	libbz2-dev libreadline-dev libsqlite3-dev wget curl llvm \
-	libncursesw5-dev xz-utils tk-dev libxml2-dev libxmlsec1-dev libffi-dev liblzma-dev git-lfs \
-	ffmpeg libsm6 libxext6 cmake libgl1 \
-	&& rm -rf /var/lib/apt/lists/* \
-	&& git lfs install
-
-RUN curl -LsSf https://astral.sh/uv/install.sh | sh
-ENV PATH="/root/.local/bin:$PATH"
-
-RUN uv venv --python 3.10 /opt/venv
-
-WORKDIR /app
-COPY requirements.txt .
-RUN . /opt/venv/bin/activate && uv pip install -r requirements.txt
-
-FROM docker.io/nvidia/cuda:12.3.2-cudnn9-runtime-ubuntu22.04
-
-ENV DEBIAN_FRONTEND=noninteractive
-
-RUN apt-get update && apt-get install -y --no-install-recommends \
-	ffmpeg libsm6 libxext6 libgl1 \
-	git-lfs \
-	&& rm -rf /var/lib/apt/lists/*
-
-RUN curl -LsSf https://astral.sh/uv/install.sh | sh
-ENV PATH="/root/.local/bin:$PATH"
-
-WORKDIR /app
-
-COPY --from=builder /opt/venv /opt/venv
-COPY . .
-
-RUN mkdir -p /app/processed_results
-RUN mkdir -p /root/.cache/torch/hub/checkpoints
-
-ENV PYTHONUNBUFFERED=1
-ENV PROCESSED_RESULTS=/app/processed_results
-
-EXPOSE 7860
-
-CMD ["uv", "run", "app.py"]

ORIGINAL_README.md

@@ -1,106 +0,0 @@
-LatentSync: Audio Conditioned Latent Diffusion Models for Lip Sync
-
-## 📖 Abstract
-We present LatentSync, an end-to-end lip sync framework based on audio conditioned latent diffusion models without any intermediate motion representation, diverging from previous diffusion-based lip sync methods based on pixel space diffusion or two-stage generation. Our framework can leverage powerful capabilities of Stable Diffusion to directly model complex audio-visual correlations. Additionally, we found that diffusion-based lip sync methods exhibit inferior temporal consistency due to inconsistency in diffusion process across different frames. We propose Temporal REPresentation Alignment (TREPA) to enhance temporal consistency while preserving lip-sync accuracy. TREPA uses temporal representations extracted by large-scale self-supervised video models to align generated frames with ground truth frames.
-
-## 🏗️ Framework
-LatentSync uses Whisper to convert melspectrogram into audio embeddings, which are then integrated into U-Net via cross-attention layers. The reference and masked frames are channel-wise concatenated with noised latents as input of U-Net. In training process, we use one-step method to get estimated clean latents from predicted noises, which are then decoded to obtain the estimated clean frames. The TREPA, LPIPS and SyncNet loss are added in the pixel space.
-
-## 🎬 Demo
-
-|  |  |
-| --- | --- |
-| __Original video__ | __Lip-synced video__ |
-| demo2_input.mp4 | demo2_output_v1.6.mp4 |
-| demo3_input.mp4 | demo3_output_v1.6.mp4 |
-| demo4_input.mp4 | demo4_output_v1.6.mp4 |
-| demo5_input.mp4 | demo5_output_v1.6.mp4 |
-| demo4_video.mp4 | demo4_output.mp4 |
-
-(Photorealistic videos are filmed by contracted models, and anime videos are from VASA-1 and EMO)
-
-## 📑 Open-source Plan
-
-- Inference code and checkpoints
-- Data processing pipeline
-- Training code
-
-## 🔧 Setting up the Environment
-Install the required packages and download the checkpoints via:
-
-```bash
-source setup_env.sh
-```
-
-If the download is successful, the checkpoints should appear as follows:
-
-```
-./checkpoints/
-|-- latentsync_unet.pt
-|-- latentsync_syncnet.pt
-|-- whisper
-|   `-- tiny.pt
-|-- auxiliary
-|   |-- 2DFAN4-cd938726ad.zip
-|   |-- i3d_torchscript.pt
-|   |-- koniq_pretrained.pkl
-|   |-- s3fd-619a316812.pth
-|   |-- sfd_face.pth
-|   |-- syncnet_v2.model
-|   |-- vgg16-397923af.pth
-|   `-- vit_g_hybrid_pt_1200e_ssv2_ft.pth
-```
-
-These already include all the checkpoints required for latentsync training and inference. If you just want to try inference, you only need to download `latentsync_unet.pt` and `tiny.pt` from our HuggingFace repo
-
-## 🚀 Inference
-Run the script for inference, which requires about 6.5 GB GPU memory.
-
-```bash
-./inference.sh
-```
-
-You can try adjusting the following inference parameters to achieve better results:
-
-- `inference_steps` [20-50]: A higher value improves visual quality but slows down the generation speed.
-- `guidance_scale` [1.0-3.0]: A higher value improves lip-sync accuracy but may cause the video distortion or jitter.
-
-## 🔄 Data Processing Pipeline
-The complete data processing pipeline includes the following steps:
-
-1. Remove the broken video files.
-2. Resample the video FPS to 25, and resample the audio to 16000 Hz.
-3. Scene detect via PySceneDetect.
-4. Split each video into 5-10 second segments.
-5. Remove videos where the face is smaller than 256 $\times$ 256, as well as videos with more than one face.
-6. Affine transform the faces according to the landmarks detected by face-alignment, then resize to 256 $\times$ 256.
-7. Remove videos with sync confidence score lower than 3, and adjust the audio-visual offset to 0.
-8. Calculate hyperIQA score, and remove videos with scores lower than 40.
-
-Run the script to execute the data processing pipeline:
-
-```bash
-./data_processing_pipeline.sh
-```
-
-You should change the parameter `input_dir` in the script to specify the data directory to be processed. The processed data will be saved in the same directory. Each step will generate a new directory to prevent the need to redo the entire pipeline in case the process is interrupted by an unexpected error.
-
-## 🏋️‍♂️ Training U-Net
-Before training, you must process the data as described above and download all the checkpoints. We released a pretrained SyncNet with 94% accuracy on VoxCeleb2 dataset for the supervision of U-Net training. Note that this SyncNet is trained on affine transformed videos, so when using or evaluating this SyncNet, you need to perform affine transformation on the video first (the code of affine transformation is included in the data processing pipeline).
-
-If all the preparations are complete, you can train the U-Net with the following script:
-
-```bash
-./train_unet.sh
-```
-
-You should change the parameters in the U-Net config file to specify the data directory, checkpoint save path, and other training hyperparameters.
-
-## 🏋️‍♂️ Training SyncNet
-In case you want to train SyncNet on your own datasets, you can run the following script. The data processing pipeline for SyncNet is the same as for U-Net.
-
-```bash
-./train_syncnet.sh
-```
-
-After `validations_steps` training, the loss charts will be saved in `train_output_dir`. They contain both the training and validation loss.

README.md

@@ -1,161 +1,10 @@
 ---
-title: LipSync - LatentSync 1.6
-emoji: 👄
-colorFrom: purple
+title: Lipsync Docker
+emoji: 🚀
+colorFrom: yellow
 colorTo: pink
-sdk: gradio
-sdk_version: 5.24.0
-python_version: "3.10"
-app_file: app.py
+sdk: docker
 pinned: false
-short_description: Lipsync video (English only) - LatentSync 1.6
 ---
 
-# OutofLipSync - LatentSync 1.6
-
-Lipsync video with custom audio (English only) using **LatentSync 1.6** from ByteDance.
-
-## Features
-
-- **Resolution**: 512x512 (LatentSync 1.6)
-- **Auto-download**: Checkpoints from `ByteDance/LatentSync-1.6`
-- **Face detection**: Automatic face detection and cropping
-- **Audio processing**: Audio separation, upsampling
-- **Multiple outputs**: Step-by-step processing visualization
-
-## HuggingFace Spaces Deployment
-
-### 1. Tạo Space mới trên HuggingFace
-
-- Vào <https://huggingface.co/new-space>
-- Chọn:
-  - **Owner**: Username của bạn
-  - **Space name**: Tên bạn muốn (ví dụ: `lipsync-demo`)
-  - **SDK**: Gradio
-  - **Hardware**: GPU (cần ít nhất 18GB VRAM cho LatentSync 1.6)
-  - **Visibility**: Public hoặc Private
-
-### 2. Đẩy code lên Space
-
-**Cách 1: Dùng Git**
-
-```bash
-git clone https://huggingface.co/spaces/YOUR_USERNAME/YOUR_SPACE_NAME
-cd YOUR_SPACE_NAME
-git remote add origin https://github.com/naicoi/OutofLipSync
-git pull origin main
-git push origin main
-```
-
-**Cách 2: Dùng HuggingFace CLI**
-
-```bash
-# Install huggingface-cli nếu chưa có
-pip install huggingface_hub
-
-# Login
-huggingface-cli login
-
-# Push code lên Space
-git push https://huggingface.co/spaces/YOUR_USERNAME/YOUR_SPACE_NAME main
-```
-
-### 3. Đợi build và deploy
-
-- HuggingFace sẽ tự động build và deploy
-- Check status ở tab "Settings" → "Build"
-- Khi build xong, app sẽ chạy tại: `https://huggingface.co/spaces/YOUR_USERNAME/YOUR_SPACE_NAME`
-
-### 4. Yêu cầu
-
-- **GPU**: Space cần có GPU (tối thiểu 18GB VRAM cho LatentSync 1.6)
-- **Runtime**: Python 3.10
-- **Disk space**: ~5GB cho checkpoints
-
-### 5. Lưu ý
-
-- Checkpoint được tải tự động từ `ByteDance/LatentSync-1.6` khi khởi động
-- Quá trình tải checkpoint có thể mất vài phút
-- Audio target chỉ hỗ trợ tiếng Anh
-
----
-
-## 🚀 Deployment
-
-### HuggingFace Spaces
-
-1. Create a Space on HuggingFace
-2. Push this repository to your Space
-3. Done! HuggingFace will automatically:
-   - Create Python environment
-   - Install dependencies from requirements.txt
-   - Start the application
-
-**Requirements:**
-
-- Hardware: A10G GPU (recommended, 24GB VRAM)
-- Python: 3.10
-
-## 💻 Local Development
-
-### Option 1: Using uv (Fast - Recommended)
-
-```bash
-# Install uv (macOS/Linux)
-curl -LsSf https://astral.sh/uv/install.sh | sh
-# Or: brew install uv
-
-# Setup and install
-./setup_local.sh
-
-# Run application
-uv run python app.py
-```
-
-**Why uv?** 10-100x faster than pip for dependency management!
-
-### Option 2: Using pip (Standard)
-
-```bash
-# Create venv
-python3 -m venv .venv
-source .venv/bin/activate  # Windows: .venv\Scripts\activate
-
-# Install dependencies
-pip install -r requirements.txt
-
-# Run application
-python app.py
-```
-
-## 📦 Dependencies
-
-- **requirements.txt**: All dependencies for application
-- Packages are installed in `.venv/` (ignored by git)
-- Git dependencies: LatentSync, FastAudioSR, tigersound, descript-audiotools
-
-## ⚠️ Important Notes
-
-### Flash-attn-3 for Local Testing
-
-The `flash-attn-3` package only provides wheels for Linux x86_64:
-
-- **HuggingFace (Linux)**: ✅ Works automatically
-- **Local (macOS)**: ❌ Will fail during installation
-
-**Workaround for local testing:**
-
-```bash
-# Comment out flash-attn-3 in requirements.txt for local testing
-# Uncomment before pushing to HuggingFace
-```
-
-### Checkpoints
-
-Checkpoints are automatically downloaded from `ByteDance/LatentSync-1.6` on startup.
-
-### Audio Language
-
-Target audio supports **English only**.
-
-Check out the configuration reference at <https://huggingface.co/docs/hub/spaces-config-reference>
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -1,189 +0,0 @@
-"""
-OutofLipSync - Lipsync Only Application
-Main Gradio UI module
-"""
-
-import os
-
-# Optimize PyTorch memory allocation to reduce fragmentation
-os.environ["PYTORCH_ALLOC_CONF"] = "expandable_segments:True"
-
-import logging
-import sys
-import shutil
-
-import gradio as gr
-import torchvision.transforms.functional as _F
-from processing import lipsync_with_audio_target
-from shared.model_manager import ModelManager
-
-
-logging.info("=" * 60)
-logging.info("APPLICATION STARTING")
-logging.info(f"Python version: {sys.version}")
-logging.info(f"Platform: {sys.platform}")
-logging.info(f"Working directory: {os.getcwd()}")
-logging.info("=" * 60)
-
-sys.modules["torchvision.transforms.functional_tensor"] = _F
-
-os.environ["PROCESSED_RESULTS"] = os.path.join(os.getcwd(), "processed_results")
-os.makedirs(os.environ["PROCESSED_RESULTS"], exist_ok=True)
-
-src = "/models"
-dst = os.path.expanduser("~/.cache/torch/hub/checkpoints")
-
-os.makedirs(dst, exist_ok=True)
-
-if os.path.exists(src):
-    for item in os.listdir(src):
-        src_path = os.path.join(src, item)
-        dst_path = os.path.join(dst, item)
-        if os.path.isfile(src_path) and not os.path.exists(dst_path):
-            shutil.copy2(src_path, dst_path)
-            print(f"Copied {item} to {dst}")
-
-print("Done copying checkpoints!")
-
-print("Loading LatentSync models...")
-manager = ModelManager.get_instance()
-manager.preload_latentsync_models()
-print("Models loaded!")
-
-
-css = """
-    #col-container {
-        margin: 0 auto;
-        max-width: 1400px;
-        padding: 2rem 1rem;
-    }
-    .header-container {
-        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
-        border-radius: 1rem;
-        padding: 2rem;
-        margin-bottom: 1.5rem;
-        box-shadow: 0 4px 20px rgba(102, 126, 234, 0.3);
-    }
-    .header-title {
-        color: white;
-        margin: 0;
-        font-size: 2.5rem;
-        font-weight: 700;
-        letter-spacing: -0.5px;
-    }
-    .header-subtitle {
-        color: rgba(255, 255, 255, 0.9);
-        margin: 0.5rem 0 0;
-        font-size: 1.1rem;
-    }
-    .card-section {
-        background: white;
-        border-radius: 1rem;
-        padding: 1.5rem;
-        box-shadow: 0 2px 12px rgba(0, 0, 0, 0.08);
-        border: 1px solid #e5e7eb;
-        height: 100%;
-        transition: all 0.3s ease;
-    }
-    .card-section:hover {
-        box-shadow: 0 4px 20px rgba(0, 0, 0, 0.12);
-    }
-    .section-header {
-        color: #1f2937;
-        font-size: 1.25rem;
-        font-weight: 600;
-        margin-bottom: 1rem;
-        display: flex;
-        align-items: center;
-        gap: 0.5rem;
-    }
-    .footer-container {
-        margin-top: 2rem;
-        padding-top: 1.5rem;
-        border-top: 1px solid #e5e7eb;
-        text-align: center;
-        color: #6b7280;
-        font-size: 0.9rem;
-    }
-    .footer-link {
-        color: #667eea;
-        text-decoration: none;
-        transition: color 0.2s ease;
-    }
-    .footer-link:hover {
-        color: #764ba2;
-    }
-    """
-
-
-def cleanup(request: gr.Request):
-    sid = request.session_hash
-    if sid:
-        print(f"{sid} left")
-        d1 = os.path.join(os.environ["PROCESSED_RESULTS"], sid)
-        shutil.rmtree(d1, ignore_errors=True)
-
-
-def start_session(request: gr.Request):
-    return request.session_hash
-
-
-with gr.Blocks(css=css) as demo:
-    session_state = gr.State()
-    demo.load(fn=start_session, outputs=[session_state])
-
-    with gr.Column(elem_id="col-container"):
-        gr.HTML(
-            """
-            <div class="header-container">
-                <h1 class="header-title">🎬 OutofLipSync</h1>
-                <p class="header-subtitle">Lipsync video with custom audio (English only)</p>
-            </div>
-            """
-        )
-
-        with gr.Row():
-            with gr.Column(scale=1):
-                with gr.Group(elem_classes="card-section"):
-                    gr.HTML('<div class="section-header">📹 Upload Video</div>')
-                    video_input = gr.Video(label="Video Source", height=400)
-
-            with gr.Column(scale=1):
-                with gr.Group(elem_classes="card-section"):
-                    gr.HTML('<div class="section-header">🎵 Upload Audio</div>')
-                    audio_input = gr.Audio(
-                        label="Target Audio (English only)", type="filepath"
-                    )
-                    quality_level = gr.Radio(
-                        choices=["Fast", "Normal", "Medium", "Best", "Super Best"],
-                        value="Normal",
-                        label="Quality",
-                    )
-                    lipsync_only_btn = gr.Button(
-                        "👄 Lipsync", variant="primary", size="lg"
-                    )
-
-        with gr.Group(elem_classes="card-section"):
-            gr.HTML('<div class="section-header">🎬 Final Output</div>')
-            final_video = gr.Video(label="Final Output", height=500)
-
-        gr.HTML(
-            """
-            <div class="footer-container">
-                <p>Made with ♥ by <a href="#" class="footer-link">LT Tech</a> • Powered by <a href="#" class="footer-link">LatentSync</a></p>
-                <p style="margin-top: 0.5rem; font-size: 0.85rem;">Version 1.0.0</p>
-            </div>
-            """
-        )
-
-    lipsync_only_btn.click(
-        fn=lipsync_with_audio_target,
-        inputs=[video_input, audio_input, session_state, quality_level],
-        outputs=[final_video],
-    )
-
-
-if __name__ == "__main__":
-    demo.unload(cleanup)
-    demo.queue()
-    demo.launch(ssr_mode=False, share=True)

audio_processing.py

@@ -1,211 +0,0 @@
-"""Audio processing utilities for OutofLipSync"""
-
-import os
-import subprocess
-from ffmpy import FFmpeg, FFRuntimeError
-
-
-def get_audio_duration(audio_path: str, max_duration: float = 30.0) -> float:
-    """Get audio file duration, raise error if exceeds max_duration
-
-    Args:
-        audio_path: Path to audio file
-        max_duration: Maximum duration in seconds (default 30)
-
-    Returns:
-        Duration in seconds
-
-    Raises:
-        ValueError: If audio duration exceeds max_duration
-    """
-    cmd = [
-        "ffprobe",
-        "-v",
-        "error",
-        "-show_entries",
-        "format=duration",
-        "-of",
-        "default=noprint_wrappers=1:nokey=1",
-        audio_path,
-    ]
-    result = subprocess.run(cmd, capture_output=True, text=True, check=True)
-    duration = float(result.stdout.strip())
-
-
-
-    return duration
-
-
-# def prepare_target_audio(audio_path: str, output_dir: str) -> tuple:
-#     """Prepare target audio for lipsync (DEPRECATED - use prepare_audio_for_lipsync instead)
-#
-#     Args:
-#         audio_path: Path to target audio
-#         output_dir: Output directory
-#
-#     Returns:
-#         (audio_16k_path, audio_upsampled_path)
-#     """
-#     audio_16k = os.path.join(output_dir, "audio_16k.wav")
-#     audio_upsampled = os.path.join(output_dir, "audio_upsampled.wav")
-#
-#     ffmpeg1 = FFmpeg(
-#         inputs={audio_path: None},
-#         outputs={
-#             audio_16k: [
-#                 "-ar",
-#                 "16000",
-#                 "-ac",
-#                 "1",
-#                 "-acodec",
-#                 "pcm_s16le",
-#                 "-loglevel",
-#                 "error",
-#                 "-y",
-#             ]
-#         },
-#     )
-#     try:
-#         ffmpeg1.run()
-#     except FFRuntimeError as e:
-#         raise Exception(f"FFmpeg failed to convert to 16k: {e}")
-#
-#     ffmpeg2 = FFmpeg(
-#         inputs={audio_16k: None},
-#         outputs={
-#             audio_upsampled: [
-#                 "-ar",
-#                 "48000",
-#                 "-ac",
-#                 "1",
-#                 "-acodec",
-#                 "pcm_s16le",
-#                 "-loglevel",
-#                 "error",
-#                 "-y",
-#             ]
-#         },
-#     )
-#     try:
-#         ffmpeg2.run()
-#     except FFRuntimeError as e:
-#         raise Exception(f"FFmpeg failed to upsample to 48k: {e}")
-#
-#     return audio_16k, audio_upsampled
-
-
-def prepare_audio_for_lipsync(audio_path: str, output_dir: str) -> str:
-    """Chuẩn bị audio 16kHz mono cho lipsync pipeline
-
-    Args:
-        audio_path: Path audio gốc
-        output_dir: Output directory
-
-    Returns:
-        Path audio 16k WAV
-    """
-    audio_16k = os.path.join(output_dir, "audio_16k.wav")
-
-    ffmpeg = FFmpeg(
-        inputs={audio_path: None},
-        outputs={
-            audio_16k: [
-                "-ar",
-                "16000",
-                "-ac",
-                "1",
-                "-acodec",
-                "pcm_s16le",
-                "-loglevel",
-                "error",
-                "-y",
-            ]
-        },
-    )
-    try:
-        ffmpeg.run()
-    except FFRuntimeError as e:
-        raise Exception(f"FFmpeg failed to convert to 16k: {e}")
-
-    return audio_16k
-
-
-def prepare_audio_for_youtube_aac(audio_path: str, output_dir: str) -> str:
-    """Chuẩn bị audio theo chuẩn YouTube (AAC)
-
-    Args:
-        audio_path: Path audio gốc
-        output_dir: Output directory
-
-    Returns:
-        Path audio YouTube (AAC)
-    """
-    from config import (
-        YOUTUBE_AUDIO_CODEC,
-        YOUTUBE_AUDIO_BITRATE,
-        YOUTUBE_AUDIO_SAMPLE_RATE,
-    )
-
-    output_path = os.path.join(output_dir, "audio_youtube.aac")
-
-    ffmpeg = FFmpeg(
-        inputs={audio_path: None},
-        outputs={
-            output_path: [
-                "-ar",
-                str(YOUTUBE_AUDIO_SAMPLE_RATE),
-                "-ac",
-                "2",
-                "-acodec",
-                YOUTUBE_AUDIO_CODEC,
-                "-b:a",
-                YOUTUBE_AUDIO_BITRATE,
-                "-loglevel",
-                "error",
-                "-y",
-            ]
-        },
-    )
-    try:
-        ffmpeg.run()
-    except FFRuntimeError as e:
-        raise Exception(f"FFmpeg failed to prepare audio for YouTube: {e}")
-
-    return output_path
-
-
-def prepare_audio_for_youtube(audio_path: str, output_dir: str) -> str:
-    """
-    Chuẩn bị audio tối ưu cho YouTube
-
-    Args:
-        audio_path: Path to audio file (WAV)
-        output_dir: Output directory
-
-    Returns:
-        Path to audio file (WAV 48kHz PCM)
-    """
-    output_path = os.path.join(output_dir, "audio_final.wav")
-
-    ffmpeg = FFmpeg(
-        inputs={audio_path: None},
-        outputs={
-            output_path: [
-                "-ar",
-                "48000",
-                "-ac",
-                "2",
-                "-acodec",
-                "pcm_s16le",
-                "-loglevel",
-                "error",
-                "-y",
-            ]
-        },
-    )
-    try:
-        ffmpeg.run()
-    except FFRuntimeError as e:
-        raise Exception(f"FFmpeg failed to prepare audio for YouTube: {e}")
-
-    return output_path

config.py

@@ -1,25 +0,0 @@
-"""Configuration constants and global settings for OutofLipSync"""
-
-import os
-
-# Models directory - from environment variable or default to /models
-MODELS_DIR = os.getenv("MODELS_DIR", "/models")
-
-# Video settings
-DEFAULT_DURATION = 10
-MIN_DURATION = 5
-MAX_DURATION = 60
-
-# Processing directory
-PROCESSED_RESULTS_DIR = "processed_results"
-
-# YouTube quality settings
-YOUTUBE_VIDEO_PRESET = "slow"
-YOUTUBE_VIDEO_CRF = 18
-YOUTUBE_VIDEO_PROFILE = "high"
-YOUTUBE_VIDEO_LEVEL = "4.2"
-YOUTUBE_VIDEO_PIX_FMT = "yuv420p"
-
-YOUTUBE_AUDIO_CODEC = "aac"
-YOUTUBE_AUDIO_BITRATE = "320k"
-YOUTUBE_AUDIO_SAMPLE_RATE = 48000

configs/inference/musetalk.yaml

@@ -1,42 +0,0 @@
-# MuseTalk V1.5 Inference Configuration
-# Simplified config for integration with OutofLipSync
-
-version: "v15"
-model_name: "musetalk"
-use_float16: true
-batch_size: 8
-fps: 25
-audio_padding_length_left: 2
-audio_padding_length_right: 2
-bbox_shift: 0
-extra_margin: 10
-left_cheek_width: 90
-right_cheek_width: 90
-parsing_mode: "jaw"
-
-# Model paths (relative to checkpoints directory)
-unet_config: "./checkpoints/musetalkV15/musetalk.json"
-unet_model: "./checkpoints/musetalkV15/unet.pth"
-vae_model: "./checkpoints/sd-vae-ft-mse"
-whisper_model: "./checkpoints/whisper-tiny"
-
-# Input paths (to be overridden programmatically)
-video_path: "data/video/input.mp4"
-audio_path: "data/audio/input.wav"
-result_dir: "./results"
-output_vid_name: null
-
-# Device settings
-gpu_id: 0
-device: "cuda"
-
-# Optional: Use saved coordinates to skip landmark detection
-use_saved_coord: false
-saved_coord: false
-
-# Audio processing
-audio_sample_rate: 16000
-
-# Video processing
-crop_size: 256
-interpolation_mode: "lanczos"

configs/scheduler_config.json

@@ -1,12 +0,0 @@
-{
-  "_class_name": "DDIMScheduler",
-  "beta_end": 0.012,
-  "beta_schedule": "scaled_linear",
-  "beta_start": 0.00085,
-  "clip_sample": false,
-  "num_train_timesteps": 1000,
-  "set_alpha_to_one": false,
-  "steps_offset": 1,
-  "trained_betas": null,
-  "skip_prk_steps": true
-}

configs/unet/stage2_512.yaml

@@ -1,99 +0,0 @@
-data:
-  syncnet_config_path: configs/syncnet/syncnet_16_pixel_attn.yaml
-  train_output_dir: debug/unet
-  train_fileslist: /mnt/bn/maliva-gen-ai-v2/chunyu.li/fileslist/data_v10_core.txt
-  train_data_dir: ""
-  audio_embeds_cache_dir: /mnt/bn/maliva-gen-ai-v2/chunyu.li/audio_cache/embeds
-  audio_mel_cache_dir: /mnt/bn/maliva-gen-ai-v2/chunyu.li/audio_cache/mel
-
-  val_video_path: assets/demo1_video.mp4
-  val_audio_path: assets/demo1_audio.wav
-  batch_size: 1
-  num_workers: 12
-  num_frames: 16
-  resolution: 512
-  mask_image_path: latentsync/utils/mask.png
-  audio_sample_rate: 16000
-  video_fps: 25
-  audio_feat_length: [2, 2]
-
-ckpt:
-  resume_ckpt_path: checkpoints/latentsync_unet.pt
-  save_ckpt_steps: 10000
-
-run:
-  pixel_space_supervise: true
-  use_syncnet: true
-  sync_loss_weight: 0.05
-  perceptual_loss_weight: 0.1
-  recon_loss_weight: 1
-  guidance_scale: 1.5
-  trepa_loss_weight: 10
-  inference_steps: 20
-  trainable_modules:
-    - motion_modules.
-    - attentions.
-  seed: 1247
-  use_mixed_noise: true
-  mixed_noise_alpha: 1
-  mixed_precision_training: true
-  enable_gradient_checkpointing: true
-  max_train_steps: 10000000
-  max_train_epochs: -1
-
-optimizer:
-  lr: 1e-5
-  scale_lr: false
-  max_grad_norm: 1.0
-  lr_scheduler: constant
-  lr_warmup_steps: 0
-
-model:
-  act_fn: silu
-  add_audio_layer: true
-  attention_head_dim: 8
-  block_out_channels: [320, 640, 1280, 1280]
-  center_input_sample: false
-  cross_attention_dim: 384
-  down_block_types:
-    [
-      "CrossAttnDownBlock3D",
-      "CrossAttnDownBlock3D",
-      "CrossAttnDownBlock3D",
-      "DownBlock3D",
-    ]
-  mid_block_type: UNetMidBlock3DCrossAttn
-  up_block_types:
-    [
-      "UpBlock3D",
-      "CrossAttnUpBlock3D",
-      "CrossAttnUpBlock3D",
-      "CrossAttnUpBlock3D",
-    ]
-  downsample_padding: 1
-  flip_sin_to_cos: true
-  freq_shift: 0
-  in_channels: 13
-  layers_per_block: 2
-  mid_block_scale_factor: 1
-  norm_eps: 1e-5
-  norm_num_groups: 32
-  out_channels: 4
-  sample_size: 64
-  resnet_time_scale_shift: default
-
-  use_motion_module: true
-  motion_module_resolutions: [1, 2, 4, 8]
-  motion_module_mid_block: false
-  motion_module_decoder_only: false
-  motion_module_type: Vanilla
-  motion_module_kwargs:
-    num_attention_heads: 8
-    num_transformer_block: 1
-    attention_block_types:
-      - Temporal_Self
-      - Temporal_Self
-    temporal_position_encoding: true
-    temporal_position_encoding_max_len: 24
-    temporal_attention_dim_div: 1
-    zero_initialize: true

docker-compose.yml

@@ -1,19 +0,0 @@
-version: '3.8'
-
-services:
-  silent-wolf:
-    build: .
-    ports:
-      - "7860:7860"
-    volumes:
-      - ./processed_results:/app/processed_results
-    deploy:
-      resources:
-        reservations:
-          devices:
-            - driver: nvidia
-              count: all
-              capabilities: [gpu]
-    environment:
-      - PYTHONUNBUFFERED=1
-      - PROCESSED_RESULTS=/app/processed_results

download_checkpoints.sh

@@ -1,18 +0,0 @@
-#!/bin/bash
-
-set -e
-
-echo "Downloading LatentSync v1.6 checkpoints..."
-
-# Create checkpoints directory
-mkdir -p checkpoints
-mkdir -p checkpoints/whisper
-
-# Download from ByteDance/LatentSync-1.6
-echo "Downloading latentsync_unet.pt..."
-huggingface-cli download ByteDance/LatentSync-1.6 latentsync_unet.pt --local-dir checkpoints
-
-echo "Downloading whisper model..."
-huggingface-cli download ByteDance/LatentSync-1.6 whisper/tiny.pt --local-dir checkpoints/whisper
-
-echo "Checkpoints downloaded successfully!"

download_musetalk_models.py

@@ -1,203 +0,0 @@
-"""Download MuseTalk V1.5 models - On-demand download manager"""
-
-from huggingface_hub import snapshot_download, hf_hub_download
-import os
-
-
-def download_musetalk_models(
-    checkpoints_dir: str = "./checkpoints", force: bool = False
-):
-    """Download MuseTalk V1.5 models if not present
-
-    Args:
-        checkpoints_dir: Directory to store checkpoints
-        force: Force re-download even if files exist
-    """
-
-    print(f"\n{'=' * 60}")
-    print(f"DOWNLOADING MUSETALK V1.5 MODELS")
-    print(f"{'=' * 60}\n")
-
-    os.makedirs(f"{checkpoints_dir}/musetalkV15", exist_ok=True)
-
-    downloaded_count = 0
-    total_count = 6
-
-    # MuseTalk V1.5 UNet and config
-    if force or not os.path.exists(f"{checkpoints_dir}/musetalkV15/unet.pth"):
-        print("  Downloading MuseTalk V1.5 UNet...")
-        try:
-            hf_hub_download(
-                repo_id="TMElyralab/MuseTalk",
-                filename="models/musetalkV15/unet.pth",
-                local_dir=checkpoints_dir,
-            )
-            downloaded_count += 1
-            print("  ✓ UNet downloaded")
-        except Exception as e:
-            print(f"  ✗ Failed to download UNet: {e}")
-    else:
-        print("  ✓ UNet already exists, skipping")
-
-    if force or not os.path.exists(f"{checkpoints_dir}/musetalkV15/musetalk.json"):
-        print("  Downloading MuseTalk V1.5 config...")
-        try:
-            hf_hub_download(
-                repo_id="TMElyralab/MuseTalk",
-                filename="models/musetalkV15/musetalk.json",
-                local_dir=checkpoints_dir,
-            )
-            downloaded_count += 1
-            print("  ✓ Config downloaded")
-        except Exception as e:
-            print(f"  ✗ Failed to download config: {e}")
-    else:
-        print("  ✓ Config already exists, skipping")
-
-    # SD-VAE-FT-MSE
-    if force or not os.path.exists(f"{checkpoints_dir}/sd-vae-ft-mse/config.json"):
-        print("  Downloading SD-VAE-FT-MSE...")
-        try:
-            snapshot_download(
-                repo_id="stabilityai/sd-vae-ft-mse",
-                local_dir=f"{checkpoints_dir}/sd-vae-ft-mse",
-            )
-            downloaded_count += 1
-            print("  ✓ SD-VAE-FT-MSE downloaded")
-        except Exception as e:
-            print(f"  ✗ Failed to download SD-VAE: {e}")
-    else:
-        print("  ✓ SD-VAE already exists, skipping")
-
-    # Whisper-Tiny
-    if force or not os.path.exists(f"{checkpoints_dir}/whisper-tiny/config.json"):
-        print("  Downloading Whisper-Tiny...")
-        try:
-            snapshot_download(
-                repo_id="openai/whisper-tiny",
-                local_dir=f"{checkpoints_dir}/whisper-tiny",
-            )
-            downloaded_count += 1
-            print("  ✓ Whisper-Tiny downloaded")
-        except Exception as e:
-            print(f"  ✗ Failed to download Whisper: {e}")
-    else:
-        print("  ✓ Whisper already exists, skipping")
-
-    # Face Parsing models
-    if force or not os.path.exists(
-        f"{checkpoints_dir}/face-parse-bisent/79999_iter.pth"
-    ):
-        print("  Downloading Face Parsing model...")
-        try:
-            hf_hub_download(
-                repo_id="TMElyralab/MuseTalk",
-                filename="models/face-parse-bisent/79999_iter.pth",
-                local_dir=f"{checkpoints_dir}/face-parse-bisent",
-            )
-            downloaded_count += 1
-            print("  ✓ Face Parsing model downloaded")
-        except Exception as e:
-            print(f"  ✗ Failed to download Face Parsing: {e}")
-    else:
-        print("  ✓ Face Parsing model already exists, skipping")
-
-    if force or not os.path.exists(
-        f"{checkpoints_dir}/face-parse-bisent/resnet18-5c106cde.pth"
-    ):
-        try:
-            hf_hub_download(
-                repo_id="TMElyralab/MuseTalk",
-                filename="models/face-parse-bisent/resnet18-5c106cde.pth",
-                local_dir=f"{checkpoints_dir}/face-parse-bisent",
-            )
-            downloaded_count += 1
-            print("  ✓ ResNet18 downloaded")
-        except Exception as e:
-            print(f"  ✗ Failed to download ResNet18: {e}")
-    else:
-        print("  ✓ ResNet18 already exists, skipping")
-
-    # DWPose models (optional - only download if mmpose is installed)
-    try:
-        import mmpose
-
-        if force or not os.path.exists(f"{checkpoints_dir}/dwpose/dw-ll_ucoco_384.pth"):
-            print("  Downloading DWPose checkpoint...")
-            try:
-                hf_hub_download(
-                    repo_id="TMElyralab/MuseTalk",
-                    filename="models/dwpose/dw-ll_ucoco_384.pth",
-                    local_dir=f"{checkpoints_dir}/dwpose",
-                )
-                downloaded_count += 1
-                print("  ✓ DWPose checkpoint downloaded")
-            except Exception as e:
-                print(f"  ✗ Failed to download DWPose: {e}")
-        else:
-            print("  ✓ DWPose checkpoint already exists, skipping")
-    except ImportError:
-        print("  ⚠ mmpose not installed, skipping DWPose download")
-
-    print(f"\n{'=' * 60}")
-    print(f"DOWNLOAD COMPLETE: {downloaded_count}/{total_count} models downloaded")
-    print(f"{'=' * 60}\n")
-
-    return downloaded_count > 0
-
-
-def check_models(checkpoints_dir: str = "./checkpoints") -> dict:
-    """Check which MuseTalk models are available
-
-    Returns:
-        Dict with status of each model
-    """
-    status = {
-        "unet": os.path.exists(f"{checkpoints_dir}/musetalkV15/unet.pth"),
-        "unet_config": os.path.exists(f"{checkpoints_dir}/musetalkV15/musetalk.json"),
-        "vae": os.path.exists(f"{checkpoints_dir}/sd-vae-ft-mse/config.json"),
-        "whisper": os.path.exists(f"{checkpoints_dir}/whisper-tiny/config.json"),
-        "face_parsing": os.path.exists(
-            f"{checkpoints_dir}/face-parse-bisent/79999_iter.pth"
-        ),
-        "resnet18": os.path.exists(
-            f"{checkpoints_dir}/face-parse-bisent/resnet18-5c106cde.pth"
-        ),
-        "dwpose": os.path.exists(f"{checkpoints_dir}/dwpose/dw-ll_ucoco_384.pth"),
-    }
-
-    missing = [k for k, v in status.items() if not v]
-    available = [k for k, v in status.items() if v]
-
-    print(f"\nModels Status:")
-    print(f"  Available: {len(available)}/{len(status)}")
-    print(f"  Missing: {', '.join(missing) if missing else 'None'}")
-
-    return status
-
-
-if __name__ == "__main__":
-    import argparse
-
-    parser = argparse.ArgumentParser(description="Download MuseTalk V1.5 models")
-    parser.add_argument(
-        "--check",
-        "-c",
-        action="store_true",
-        help="Check model status without downloading",
-    )
-    parser.add_argument(
-        "--force", "-f", action="store_true", help="Force re-download all models"
-    )
-    parser.add_argument(
-        "--checkpoint-dir",
-        default="./checkpoints",
-        help="Directory to store checkpoints",
-    )
-
-    args = parser.parse_args()
-
-    if args.check:
-        check_models(args.checkpoint_dir)
-    else:
-        download_musetalk_models(args.checkpoint_dir, force=args.force)

eval/detectors/README.md

@@ -1,3 +0,0 @@
-# Face detector
-
-This face detector is adapted from `https://github.com/cs-giung/face-detection-pytorch`.

eval/detectors/__init__.py

@@ -1 +0,0 @@
-from .s3fd import S3FD

eval/detectors/s3fd/__init__.py

@@ -1,63 +0,0 @@
-import time
-import numpy as np
-import cv2
-import torch
-from torchvision import transforms
-from .nets import S3FDNet
-from .box_utils import nms_
-from latentsync.utils.util import check_model_and_download
-
-PATH_WEIGHT = "checkpoints/auxiliary/sfd_face.pth"
-img_mean = np.array([104.0, 117.0, 123.0])[:, np.newaxis, np.newaxis].astype("float32")
-
-
-class S3FD:
-
-    def __init__(self, device="cuda"):
-
-        tstamp = time.time()
-        self.device = device
-
-        print("[S3FD] loading with", self.device)
-        self.net = S3FDNet(device=self.device).to(self.device)
-        check_model_and_download(PATH_WEIGHT)
-        state_dict = torch.load(PATH_WEIGHT, map_location=self.device, weights_only=True)
-        self.net.load_state_dict(state_dict)
-        self.net.eval()
-        print("[S3FD] finished loading (%.4f sec)" % (time.time() - tstamp))
-
-    def detect_faces(self, image, conf_th=0.8, scales=[1]):
-
-        w, h = image.shape[1], image.shape[0]
-
-        bboxes = np.empty(shape=(0, 5))
-
-        with torch.no_grad():
-            for s in scales:
-                scaled_img = cv2.resize(image, dsize=(0, 0), fx=s, fy=s, interpolation=cv2.INTER_LINEAR)
-
-                scaled_img = np.swapaxes(scaled_img, 1, 2)
-                scaled_img = np.swapaxes(scaled_img, 1, 0)
-                scaled_img = scaled_img[[2, 1, 0], :, :]
-                scaled_img = scaled_img.astype("float32")
-                scaled_img -= img_mean
-                scaled_img = scaled_img[[2, 1, 0], :, :]
-                x = torch.from_numpy(scaled_img).unsqueeze(0).to(self.device)
-                y = self.net(x)
-
-                detections = y.data
-                scale = torch.Tensor([w, h, w, h])
-
-                for i in range(detections.size(1)):
-                    j = 0
-                    while detections[0, i, j, 0] > conf_th:
-                        score = detections[0, i, j, 0]
-                        pt = (detections[0, i, j, 1:] * scale).cpu().numpy()
-                        bbox = (pt[0], pt[1], pt[2], pt[3], score)
-                        bboxes = np.vstack((bboxes, bbox))
-                        j += 1
-
-            keep = nms_(bboxes, 0.1)
-            bboxes = bboxes[keep]
-
-        return bboxes

eval/detectors/s3fd/box_utils.py

@@ -1,221 +0,0 @@
-import numpy as np
-from itertools import product as product
-import torch
-from torch.autograd import Function
-import warnings
-
-
-def nms_(dets, thresh):
-    """
-    Courtesy of Ross Girshick
-    [https://github.com/rbgirshick/py-faster-rcnn/blob/master/lib/nms/py_cpu_nms.py]
-    """
-    x1 = dets[:, 0]
-    y1 = dets[:, 1]
-    x2 = dets[:, 2]
-    y2 = dets[:, 3]
-    scores = dets[:, 4]
-
-    areas = (x2 - x1) * (y2 - y1)
-    order = scores.argsort()[::-1]
-
-    keep = []
-    while order.size > 0:
-        i = order[0]
-        keep.append(int(i))
-        xx1 = np.maximum(x1[i], x1[order[1:]])
-        yy1 = np.maximum(y1[i], y1[order[1:]])
-        xx2 = np.minimum(x2[i], x2[order[1:]])
-        yy2 = np.minimum(y2[i], y2[order[1:]])
-
-        w = np.maximum(0.0, xx2 - xx1)
-        h = np.maximum(0.0, yy2 - yy1)
-        inter = w * h
-        ovr = inter / (areas[i] + areas[order[1:]] - inter)
-
-        inds = np.where(ovr <= thresh)[0]
-        order = order[inds + 1]
-
-    return np.array(keep).astype(np.int32)
-
-
-def decode(loc, priors, variances):
-    """Decode locations from predictions using priors to undo
-    the encoding we did for offset regression at train time.
-    Args:
-        loc (tensor): location predictions for loc layers,
-            Shape: [num_priors,4]
-        priors (tensor): Prior boxes in center-offset form.
-            Shape: [num_priors,4].
-        variances: (list[float]) Variances of priorboxes
-    Return:
-        decoded bounding box predictions
-    """
-
-    boxes = torch.cat((
-        priors[:, :2] + loc[:, :2] * variances[0] * priors[:, 2:],
-        priors[:, 2:] * torch.exp(loc[:, 2:] * variances[1])), 1)
-    boxes[:, :2] -= boxes[:, 2:] / 2
-    boxes[:, 2:] += boxes[:, :2]
-    return boxes
-
-
-def nms(boxes, scores, overlap=0.5, top_k=200):
-    """Apply non-maximum suppression at test time to avoid detecting too many
-    overlapping bounding boxes for a given object.
-    Args:
-        boxes: (tensor) The location preds for the img, Shape: [num_priors,4].
-        scores: (tensor) The class predscores for the img, Shape:[num_priors].
-        overlap: (float) The overlap thresh for suppressing unnecessary boxes.
-        top_k: (int) The Maximum number of box preds to consider.
-    Return:
-        The indices of the kept boxes with respect to num_priors.
-    """
-
-    keep = scores.new(scores.size(0)).zero_().long()
-    if boxes.numel() == 0:
-        return keep, 0
-    x1 = boxes[:, 0]
-    y1 = boxes[:, 1]
-    x2 = boxes[:, 2]
-    y2 = boxes[:, 3]
-    area = torch.mul(x2 - x1, y2 - y1)
-    v, idx = scores.sort(0)  # sort in ascending order
-    # I = I[v >= 0.01]
-    idx = idx[-top_k:]  # indices of the top-k largest vals
-    xx1 = boxes.new()
-    yy1 = boxes.new()
-    xx2 = boxes.new()
-    yy2 = boxes.new()
-    w = boxes.new()
-    h = boxes.new()
-
-    # keep = torch.Tensor()
-    count = 0
-    while idx.numel() > 0:
-        i = idx[-1]  # index of current largest val
-        # keep.append(i)
-        keep[count] = i
-        count += 1
-        if idx.size(0) == 1:
-            break
-        idx = idx[:-1]  # remove kept element from view
-        # load bboxes of next highest vals
-        with warnings.catch_warnings():
-            # Ignore UserWarning within this block
-            warnings.simplefilter("ignore", category=UserWarning)
-            torch.index_select(x1, 0, idx, out=xx1)
-            torch.index_select(y1, 0, idx, out=yy1)
-            torch.index_select(x2, 0, idx, out=xx2)
-            torch.index_select(y2, 0, idx, out=yy2)
-        # store element-wise max with next highest score
-        xx1 = torch.clamp(xx1, min=x1[i])
-        yy1 = torch.clamp(yy1, min=y1[i])
-        xx2 = torch.clamp(xx2, max=x2[i])
-        yy2 = torch.clamp(yy2, max=y2[i])
-        w.resize_as_(xx2)
-        h.resize_as_(yy2)
-        w = xx2 - xx1
-        h = yy2 - yy1
-        # check sizes of xx1 and xx2.. after each iteration
-        w = torch.clamp(w, min=0.0)
-        h = torch.clamp(h, min=0.0)
-        inter = w * h
-        # IoU = i / (area(a) + area(b) - i)
-        rem_areas = torch.index_select(area, 0, idx)  # load remaining areas)
-        union = (rem_areas - inter) + area[i]
-        IoU = inter / union  # store result in iou
-        # keep only elements with an IoU <= overlap
-        idx = idx[IoU.le(overlap)]
-    return keep, count
-
-
-class Detect(object):
-
-    def __init__(self, num_classes=2,
-                    top_k=750, nms_thresh=0.3, conf_thresh=0.05,
-                    variance=[0.1, 0.2], nms_top_k=5000):
-        
-        self.num_classes = num_classes
-        self.top_k = top_k
-        self.nms_thresh = nms_thresh
-        self.conf_thresh = conf_thresh
-        self.variance = variance
-        self.nms_top_k = nms_top_k
-
-    def forward(self, loc_data, conf_data, prior_data):
-
-        num = loc_data.size(0)
-        num_priors = prior_data.size(0)
-
-        conf_preds = conf_data.view(num, num_priors, self.num_classes).transpose(2, 1)
-        batch_priors = prior_data.view(-1, num_priors, 4).expand(num, num_priors, 4)
-        batch_priors = batch_priors.contiguous().view(-1, 4)
-
-        decoded_boxes = decode(loc_data.view(-1, 4), batch_priors, self.variance)
-        decoded_boxes = decoded_boxes.view(num, num_priors, 4)
-
-        output = torch.zeros(num, self.num_classes, self.top_k, 5)
-
-        for i in range(num):
-            boxes = decoded_boxes[i].clone()
-            conf_scores = conf_preds[i].clone()
-
-            for cl in range(1, self.num_classes):
-                c_mask = conf_scores[cl].gt(self.conf_thresh)
-                scores = conf_scores[cl][c_mask]
-                
-                if scores.dim() == 0:
-                    continue
-                l_mask = c_mask.unsqueeze(1).expand_as(boxes)
-                boxes_ = boxes[l_mask].view(-1, 4)
-                ids, count = nms(boxes_, scores, self.nms_thresh, self.nms_top_k)
-                count = count if count < self.top_k else self.top_k
-
-                output[i, cl, :count] = torch.cat((scores[ids[:count]].unsqueeze(1), boxes_[ids[:count]]), 1)
-
-        return output
-
-
-class PriorBox(object):
-
-    def __init__(self, input_size, feature_maps,
-                    variance=[0.1, 0.2],
-                    min_sizes=[16, 32, 64, 128, 256, 512],
-                    steps=[4, 8, 16, 32, 64, 128],
-                    clip=False):
-
-        super(PriorBox, self).__init__()
-
-        self.imh = input_size[0]
-        self.imw = input_size[1]
-        self.feature_maps = feature_maps
-
-        self.variance = variance
-        self.min_sizes = min_sizes
-        self.steps = steps
-        self.clip = clip
-
-    def forward(self):
-        mean = []
-        for k, fmap in enumerate(self.feature_maps):
-            feath = fmap[0]
-            featw = fmap[1]
-            for i, j in product(range(feath), range(featw)):
-                f_kw = self.imw / self.steps[k]
-                f_kh = self.imh / self.steps[k]
-
-                cx = (j + 0.5) / f_kw
-                cy = (i + 0.5) / f_kh
-
-                s_kw = self.min_sizes[k] / self.imw
-                s_kh = self.min_sizes[k] / self.imh
-
-                mean += [cx, cy, s_kw, s_kh]
-
-        output = torch.FloatTensor(mean).view(-1, 4)
-        
-        if self.clip:
-            output.clamp_(max=1, min=0)
-        
-        return output

eval/detectors/s3fd/nets.py

@@ -1,174 +0,0 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import torch.nn.init as init
-from .box_utils import Detect, PriorBox
-
-
-class L2Norm(nn.Module):
-
-    def __init__(self, n_channels, scale):
-        super(L2Norm, self).__init__()
-        self.n_channels = n_channels
-        self.gamma = scale or None
-        self.eps = 1e-10
-        self.weight = nn.Parameter(torch.Tensor(self.n_channels))
-        self.reset_parameters()
-
-    def reset_parameters(self):
-        init.constant_(self.weight, self.gamma)
-
-    def forward(self, x):
-        norm = x.pow(2).sum(dim=1, keepdim=True).sqrt() + self.eps
-        x = torch.div(x, norm)
-        out = self.weight.unsqueeze(0).unsqueeze(2).unsqueeze(3).expand_as(x) * x
-        return out
-
-
-class S3FDNet(nn.Module):
-
-    def __init__(self, device='cuda'):
-        super(S3FDNet, self).__init__()
-        self.device = device
-
-        self.vgg = nn.ModuleList([
-            nn.Conv2d(3, 64, 3, 1, padding=1),
-            nn.ReLU(inplace=True),
-            nn.Conv2d(64, 64, 3, 1, padding=1),
-            nn.ReLU(inplace=True),
-            nn.MaxPool2d(2, 2),
-
-            nn.Conv2d(64, 128, 3, 1, padding=1),
-            nn.ReLU(inplace=True),
-            nn.Conv2d(128, 128, 3, 1, padding=1),
-            nn.ReLU(inplace=True),
-            nn.MaxPool2d(2, 2),
-            
-            nn.Conv2d(128, 256, 3, 1, padding=1),
-            nn.ReLU(inplace=True),
-            nn.Conv2d(256, 256, 3, 1, padding=1),
-            nn.ReLU(inplace=True),
-            nn.Conv2d(256, 256, 3, 1, padding=1),
-            nn.ReLU(inplace=True),
-            nn.MaxPool2d(2, 2, ceil_mode=True),
-            
-            nn.Conv2d(256, 512, 3, 1, padding=1),
-            nn.ReLU(inplace=True),
-            nn.Conv2d(512, 512, 3, 1, padding=1),
-            nn.ReLU(inplace=True),
-            nn.Conv2d(512, 512, 3, 1, padding=1),
-            nn.ReLU(inplace=True),
-            nn.MaxPool2d(2, 2),
-
-            nn.Conv2d(512, 512, 3, 1, padding=1),
-            nn.ReLU(inplace=True),
-            nn.Conv2d(512, 512, 3, 1, padding=1),
-            nn.ReLU(inplace=True),
-            nn.Conv2d(512, 512, 3, 1, padding=1),
-            nn.ReLU(inplace=True),
-            nn.MaxPool2d(2, 2),
-
-            nn.Conv2d(512, 1024, 3, 1, padding=6, dilation=6),
-            nn.ReLU(inplace=True),
-            nn.Conv2d(1024, 1024, 1, 1),
-            nn.ReLU(inplace=True),
-        ])
-
-        self.L2Norm3_3 = L2Norm(256, 10)
-        self.L2Norm4_3 = L2Norm(512, 8)
-        self.L2Norm5_3 = L2Norm(512, 5)
-
-        self.extras = nn.ModuleList([
-            nn.Conv2d(1024, 256, 1, 1),
-            nn.Conv2d(256, 512, 3, 2, padding=1),
-            nn.Conv2d(512, 128, 1, 1),
-            nn.Conv2d(128, 256, 3, 2, padding=1),
-        ])
-        
-        self.loc = nn.ModuleList([
-            nn.Conv2d(256, 4, 3, 1, padding=1),
-            nn.Conv2d(512, 4, 3, 1, padding=1),
-            nn.Conv2d(512, 4, 3, 1, padding=1),
-            nn.Conv2d(1024, 4, 3, 1, padding=1),
-            nn.Conv2d(512, 4, 3, 1, padding=1),
-            nn.Conv2d(256, 4, 3, 1, padding=1),
-        ])
-
-        self.conf = nn.ModuleList([
-            nn.Conv2d(256, 4, 3, 1, padding=1),
-            nn.Conv2d(512, 2, 3, 1, padding=1),
-            nn.Conv2d(512, 2, 3, 1, padding=1),
-            nn.Conv2d(1024, 2, 3, 1, padding=1),
-            nn.Conv2d(512, 2, 3, 1, padding=1),
-            nn.Conv2d(256, 2, 3, 1, padding=1),
-        ])
-
-        self.softmax = nn.Softmax(dim=-1)
-        self.detect = Detect()
-
-    def forward(self, x):
-        size = x.size()[2:]
-        sources = list()
-        loc = list()
-        conf = list()
-
-        for k in range(16):
-            x = self.vgg[k](x)
-        s = self.L2Norm3_3(x)
-        sources.append(s)
-
-        for k in range(16, 23):
-            x = self.vgg[k](x)
-        s = self.L2Norm4_3(x)
-        sources.append(s)
-
-        for k in range(23, 30):
-            x = self.vgg[k](x)
-        s = self.L2Norm5_3(x)
-        sources.append(s)
-
-        for k in range(30, len(self.vgg)):
-            x = self.vgg[k](x)
-        sources.append(x)
-        
-        # apply extra layers and cache source layer outputs
-        for k, v in enumerate(self.extras):
-            x = F.relu(v(x), inplace=True)
-            if k % 2 == 1:
-                sources.append(x)
-
-        # apply multibox head to source layers
-        loc_x = self.loc[0](sources[0])
-        conf_x = self.conf[0](sources[0])
-
-        max_conf, _ = torch.max(conf_x[:, 0:3, :, :], dim=1, keepdim=True)
-        conf_x = torch.cat((max_conf, conf_x[:, 3:, :, :]), dim=1)
-
-        loc.append(loc_x.permute(0, 2, 3, 1).contiguous())
-        conf.append(conf_x.permute(0, 2, 3, 1).contiguous())
-
-        for i in range(1, len(sources)):
-            x = sources[i]
-            conf.append(self.conf[i](x).permute(0, 2, 3, 1).contiguous())
-            loc.append(self.loc[i](x).permute(0, 2, 3, 1).contiguous())
-
-        features_maps = []
-        for i in range(len(loc)):
-            feat = []
-            feat += [loc[i].size(1), loc[i].size(2)]
-            features_maps += [feat]
-
-        loc = torch.cat([o.view(o.size(0), -1) for o in loc], 1)
-        conf = torch.cat([o.view(o.size(0), -1) for o in conf], 1)
-
-        with torch.no_grad():
-            self.priorbox = PriorBox(size, features_maps)
-            self.priors = self.priorbox.forward()
-
-        output = self.detect.forward(
-            loc.view(loc.size(0), -1, 4),
-            self.softmax(conf.view(conf.size(0), -1, 2)),
-            self.priors.type(type(x.data)).to(self.device)
-        )
-
-        return output

eval/draw_syncnet_lines.py

@@ -1,64 +0,0 @@
-# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import torch
-import matplotlib.pyplot as plt
-
-
-class Chart:
-    def __init__(self):
-        self.loss_list = []
-
-    def add_ckpt(self, ckpt_path, line_name):
-        ckpt = torch.load(ckpt_path, map_location="cpu")
-        train_step_list = ckpt["train_step_list"]
-        train_loss_list = ckpt["train_loss_list"]
-        val_step_list = ckpt["val_step_list"]
-        val_loss_list = ckpt["val_loss_list"]
-        self.loss_list.append((line_name, train_step_list, train_loss_list, val_step_list, val_loss_list))
-
-    def draw(self, save_path, plot_val=True):
-        # Global settings
-        plt.rcParams["font.size"] = 14
-        plt.rcParams["font.family"] = "serif"
-        plt.rcParams["font.sans-serif"] = ["Arial", "DejaVu Sans", "Lucida Grande"]
-        plt.rcParams["font.serif"] = ["Times New Roman", "DejaVu Serif"]
-
-        # Creating the plot
-        plt.figure(figsize=(7.766, 4.8))  # Golden ratio
-        for loss in self.loss_list:
-            if plot_val:
-                (line,) = plt.plot(loss[1], loss[2], label=loss[0], linewidth=0.5, alpha=0.5)
-                line_color = line.get_color()
-                plt.plot(loss[3], loss[4], linewidth=1.5, color=line_color)
-            else:
-                plt.plot(loss[1], loss[2], label=loss[0], linewidth=1)
-        plt.xlabel("Step")
-        plt.ylabel("Loss")
-        legend = plt.legend()
-        # legend = plt.legend(loc='upper right', bbox_to_anchor=(1, 0.82))
-
-        # Adjust the linewidth of legend
-        for line in legend.get_lines():
-            line.set_linewidth(2)
-
-        plt.savefig(save_path, transparent=True)
-        plt.close()
-
-
-if __name__ == "__main__":
-    chart = Chart()
-    chart.add_ckpt("output/syncnet/train-2024_10_28-23:16:40/checkpoints/checkpoint-20000.pt", "Wav2Lip SyncNet")
-    chart.add_ckpt("output/syncnet/train-2024_10_29-20:13:43/checkpoints/checkpoint-20000.pt", "StableSyncNet")
-    chart.draw("ablation.pdf", plot_val=True)

eval/eval_fvd.py

@@ -1,98 +0,0 @@
-# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import mediapipe as mp
-import cv2
-from decord import VideoReader
-import os
-import numpy as np
-import torch
-import tqdm
-from eval.fvd import compute_our_fvd
-
-
-class FVD:
-    def __init__(self, resolution=(224, 224)):
-        self.face_detector = mp.solutions.face_detection.FaceDetection(model_selection=0, min_detection_confidence=0.5)
-        self.resolution = resolution
-
-    def detect_face(self, image):
-        height, width = image.shape[:2]
-        # Process the image and detect faces.
-        results = self.face_detector.process(image)
-
-        if not results.detections:  # Face not detected
-            raise RuntimeError("Face not detected")
-
-        detection = results.detections[0]  # Only use the first face in the image
-        bounding_box = detection.location_data.relative_bounding_box
-        xmin = int(bounding_box.xmin * width)
-        ymin = int(bounding_box.ymin * height)
-        face_width = int(bounding_box.width * width)
-        face_height = int(bounding_box.height * height)
-
-        # Crop the image to the bounding box.
-        xmin = max(0, xmin)
-        ymin = max(0, ymin)
-        xmax = min(width, xmin + face_width)
-        ymax = min(height, ymin + face_height)
-        image = image[ymin:ymax, xmin:xmax]
-
-        return image
-
-    def detect_video(self, video_path):
-        vr = VideoReader(video_path)
-        video_frames = vr[20:36].asnumpy()
-        vr.seek(0)  # avoid memory leak
-        faces = []
-        for frame in video_frames:
-            face = self.detect_face(frame)
-            face = cv2.resize(face, (self.resolution[1], self.resolution[0]), interpolation=cv2.INTER_AREA)
-            faces.append(face)
-
-        if len(faces) != 16:
-            return RuntimeError("Insufficient consecutive frames of faces (less than 16).")
-        faces = np.stack(faces, axis=0)  # (f, h, w, c)
-        faces = torch.from_numpy(faces)
-        return faces
-
-    def detect_videos(self, videos_dir: str):
-        videos_list = []
-
-        if videos_dir.endswith(".mp4"):
-            video_faces = self.detect_video(videos_dir)
-            videos_list.append(video_faces)
-        else:
-            for file in tqdm.tqdm(os.listdir(videos_dir)):
-                if file.endswith(".mp4"):
-                    video_path = os.path.join(videos_dir, file)
-                    video_faces = self.detect_video(video_path)
-                    videos_list.append(video_faces)
-
-        videos_list = torch.stack(videos_list) / 255.0
-        return videos_list
-
-
-def eval_fvd(real_videos_dir: str, fake_videos_dir: str):
-    fvd = FVD()
-    real_videos = fvd.detect_videos(real_videos_dir)
-    fake_videos = fvd.detect_videos(fake_videos_dir)
-    fvd_value = compute_our_fvd(real_videos, fake_videos, device="cpu")
-    print(f"FVD: {fvd_value:.3f}")
-
-
-if __name__ == "__main__":
-    real_videos_dir = "dir1"
-    fake_videos_dir = "dir2"
-    eval_fvd(real_videos_dir, fake_videos_dir)

eval/eval_sync_conf.py

@@ -1,77 +0,0 @@
-# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import argparse
-import os
-import tqdm
-from statistics import fmean
-from eval.syncnet import SyncNetEval
-from eval.syncnet_detect import SyncNetDetector
-from latentsync.utils.util import red_text
-import torch
-
-
-def syncnet_eval(syncnet, syncnet_detector, video_path, temp_dir, detect_results_dir="detect_results"):
-    syncnet_detector(video_path=video_path, min_track=50)
-    crop_videos = os.listdir(os.path.join(detect_results_dir, "crop"))
-    if crop_videos == []:
-        raise Exception(red_text(f"Face not detected in {video_path}"))
-    av_offset_list = []
-    conf_list = []
-    for video in crop_videos:
-        av_offset, _, conf = syncnet.evaluate(
-            video_path=os.path.join(detect_results_dir, "crop", video), temp_dir=temp_dir
-        )
-        av_offset_list.append(av_offset)
-        conf_list.append(conf)
-    av_offset = int(fmean(av_offset_list))
-    conf = fmean(conf_list)
-    print(f"Input video: {video_path}\nSyncNet confidence: {conf:.2f}\nAV offset: {av_offset}")
-    return av_offset, conf
-
-
-def main():
-    parser = argparse.ArgumentParser(description="SyncNet")
-    parser.add_argument("--initial_model", type=str, default="checkpoints/auxiliary/syncnet_v2.model", help="")
-    parser.add_argument("--video_path", type=str, default=None, help="")
-    parser.add_argument("--videos_dir", type=str, default="/root/processed")
-    parser.add_argument("--temp_dir", type=str, default="temp", help="")
-
-    args = parser.parse_args()
-
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-
-    syncnet = SyncNetEval(device=device)
-    syncnet.loadParameters(args.initial_model)
-
-    syncnet_detector = SyncNetDetector(device=device, detect_results_dir="detect_results")
-
-    if args.video_path is not None:
-        syncnet_eval(syncnet, syncnet_detector, args.video_path, args.temp_dir)
-    else:
-        sync_conf_list = []
-        video_names = sorted([f for f in os.listdir(args.videos_dir) if f.endswith(".mp4")])
-        for video_name in tqdm.tqdm(video_names):
-            try:
-                _, conf = syncnet_eval(
-                    syncnet, syncnet_detector, os.path.join(args.videos_dir, video_name), args.temp_dir
-                )
-                sync_conf_list.append(conf)
-            except Exception as e:
-                print(e)
-        print(f"The average sync confidence is {fmean(sync_conf_list):.02f}")
-
-
-if __name__ == "__main__":
-    main()

eval/eval_sync_conf.sh

@@ -1,2 +0,0 @@
-#!/bin/bash
-python -m eval.eval_sync_conf --video_path "video_out.mp4"

eval/eval_syncnet_acc.py

@@ -1,137 +0,0 @@
-# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import argparse
-import os
-import sys
-from tqdm.auto import tqdm
-import torch
-import torch.nn as nn
-from einops import rearrange
-from latentsync.models.stable_syncnet import StableSyncNet
-from latentsync.data.syncnet_dataset import SyncNetDataset
-from diffusers import AutoencoderKL
-from omegaconf import OmegaConf
-from accelerate.utils import set_seed
-
-sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-from config import MODELS_DIR
-
-
-def main(config):
-    set_seed(config.run.seed)
-
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-
-    if config.data.latent_space:
-        vae = AutoencoderKL.from_pretrained(
-            "runwayml/stable-diffusion-inpainting",
-            subfolder="vae",
-            revision="fp16",
-            torch_dtype=torch.float16,
-            cache_dir=MODELS_DIR,
-        )
-        vae.requires_grad_(False)
-        vae.to(device)
-
-    # Dataset and Dataloader setup
-    dataset = SyncNetDataset(
-        config.data.val_data_dir, config.data.val_fileslist, config
-    )
-
-    test_dataloader = torch.utils.data.DataLoader(
-        dataset,
-        batch_size=config.data.batch_size,
-        shuffle=False,
-        num_workers=config.data.num_workers,
-        drop_last=False,
-        worker_init_fn=dataset.worker_init_fn,
-    )
-
-    # Model
-    syncnet = StableSyncNet(OmegaConf.to_container(config.model)).to(device)
-
-    print(f"Load checkpoint from: {config.ckpt.inference_ckpt_path}")
-    checkpoint = torch.load(
-        config.ckpt.inference_ckpt_path, map_location=device, weights_only=True
-    )
-
-    syncnet.load_state_dict(checkpoint["state_dict"])
-    syncnet.to(dtype=torch.float16)
-    syncnet.requires_grad_(False)
-    syncnet.eval()
-
-    global_step = 0
-    num_val_batches = config.data.num_val_samples // config.data.batch_size
-    progress_bar = tqdm(range(0, num_val_batches), initial=0, desc="Testing accuracy")
-
-    num_correct_preds = 0
-    num_total_preds = 0
-
-    while True:
-        for step, batch in enumerate(test_dataloader):
-            ### >>>> Test >>>> ###
-
-            frames = batch["frames"].to(device, dtype=torch.float16)
-            audio_samples = batch["audio_samples"].to(device, dtype=torch.float16)
-            y = batch["y"].to(device, dtype=torch.float16).squeeze(1)
-
-            if config.data.latent_space:
-                frames = rearrange(frames, "b f c h w -> (b f) c h w")
-
-                with torch.no_grad():
-                    frames = vae.encode(frames).latent_dist.sample() * 0.18215
-
-                frames = rearrange(
-                    frames, "(b f) c h w -> b (f c) h w", f=config.data.num_frames
-                )
-            else:
-                frames = rearrange(frames, "b f c h w -> b (f c) h w")
-
-            if config.data.lower_half:
-                height = frames.shape[2]
-                frames = frames[:, :, height // 2 :, :]
-
-            with torch.no_grad():
-                vision_embeds, audio_embeds = syncnet(frames, audio_samples)
-
-            sims = nn.functional.cosine_similarity(vision_embeds, audio_embeds)
-
-            preds = (sims > 0.5).to(dtype=torch.float16)
-            num_correct_preds += (preds == y).sum().item()
-            num_total_preds += len(sims)
-
-            progress_bar.update(1)
-            global_step += 1
-
-            if global_step >= num_val_batches:
-                progress_bar.close()
-                print(
-                    f"SyncNet Accuracy: {num_correct_preds / num_total_preds * 100:.2f}%"
-                )
-                return
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="Code to test the accuracy of SyncNet")
-
-    parser.add_argument(
-        "--config_path", type=str, default="configs/syncnet/syncnet_16_latent.yaml"
-    )
-    args = parser.parse_args()
-
-    # Load a configuration file
-    config = OmegaConf.load(args.config_path)
-
-    main(config)

eval/eval_syncnet_acc.sh

@@ -1,3 +0,0 @@
-#!/bin/bash
-
-python -m eval.eval_syncnet_acc --config_path "configs/syncnet/syncnet_16_pixel_attn.yaml"

eval/fvd.py

@@ -1,58 +0,0 @@
-# Adapted from https://github.com/universome/fvd-comparison/blob/master/our_fvd.py
-
-from typing import Tuple
-import scipy
-import numpy as np
-import torch
-from latentsync.utils.util import check_model_and_download
-
-
-def compute_fvd(feats_fake: np.ndarray, feats_real: np.ndarray) -> float:
-    mu_gen, sigma_gen = compute_stats(feats_fake)
-    mu_real, sigma_real = compute_stats(feats_real)
-
-    m = np.square(mu_gen - mu_real).sum()
-    s, _ = scipy.linalg.sqrtm(np.dot(sigma_gen, sigma_real), disp=False)  # pylint: disable=no-member
-    fid = np.real(m + np.trace(sigma_gen + sigma_real - s * 2))
-
-    return float(fid)
-
-
-def compute_stats(feats: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
-    mu = feats.mean(axis=0)  # [d]
-    sigma = np.cov(feats, rowvar=False)  # [d, d]
-
-    return mu, sigma
-
-
-@torch.no_grad()
-def compute_our_fvd(videos_fake: np.ndarray, videos_real: np.ndarray, device: str = "cuda") -> float:
-    i3d_path = "checkpoints/auxiliary/i3d_torchscript.pt"
-    check_model_and_download(i3d_path)
-    i3d_kwargs = dict(
-        rescale=False, resize=False, return_features=True
-    )  # Return raw features before the softmax layer.
-
-    with open(i3d_path, "rb") as f:
-        i3d_model = torch.jit.load(f).eval().to(device)
-
-    videos_fake = videos_fake.permute(0, 4, 1, 2, 3).to(device)
-    videos_real = videos_real.permute(0, 4, 1, 2, 3).to(device)
-
-    feats_fake = i3d_model(videos_fake, **i3d_kwargs).cpu().numpy()
-    feats_real = i3d_model(videos_real, **i3d_kwargs).cpu().numpy()
-
-    return compute_fvd(feats_fake, feats_real)
-
-
-def main():
-    # input shape: (b, f, h, w, c)
-    videos_fake = torch.rand(10, 16, 224, 224, 3)
-    videos_real = torch.rand(10, 16, 224, 224, 3)
-
-    our_fvd_result = compute_our_fvd(videos_fake, videos_real)
-    print(f"[FVD scores] Ours: {our_fvd_result}")
-
-
-if __name__ == "__main__":
-    main()

eval/hyper_iqa.py

@@ -1,343 +0,0 @@
-# Adapted from https://github.com/SSL92/hyperIQA/blob/master/models.py
-
-import torch as torch
-import torch.nn as nn
-from torch.nn import functional as F
-from torch.nn import init
-import math
-import torch.utils.model_zoo as model_zoo
-
-model_urls = {
-    'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
-    'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
-    'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
-    'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
-    'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
-}
-
-
-class HyperNet(nn.Module):
-    """
-    Hyper network for learning perceptual rules.
-
-    Args:
-        lda_out_channels: local distortion aware module output size.
-        hyper_in_channels: input feature channels for hyper network.
-        target_in_size: input vector size for target network.
-        target_fc(i)_size: fully connection layer size of target network.
-        feature_size: input feature map width/height for hyper network.
-
-    Note:
-        For size match, input args must satisfy: 'target_fc(i)_size * target_fc(i+1)_size' is divisible by 'feature_size ^ 2'.
-
-    """
-    def __init__(self, lda_out_channels, hyper_in_channels, target_in_size, target_fc1_size, target_fc2_size, target_fc3_size, target_fc4_size, feature_size):
-        super(HyperNet, self).__init__()
-
-        self.hyperInChn = hyper_in_channels
-        self.target_in_size = target_in_size
-        self.f1 = target_fc1_size
-        self.f2 = target_fc2_size
-        self.f3 = target_fc3_size
-        self.f4 = target_fc4_size
-        self.feature_size = feature_size
-
-        self.res = resnet50_backbone(lda_out_channels, target_in_size, pretrained=True)
-
-        self.pool = nn.AdaptiveAvgPool2d((1, 1))
-
-        # Conv layers for resnet output features
-        self.conv1 = nn.Sequential(
-            nn.Conv2d(2048, 1024, 1, padding=(0, 0)),
-            nn.ReLU(inplace=True),
-            nn.Conv2d(1024, 512, 1, padding=(0, 0)),
-            nn.ReLU(inplace=True),
-            nn.Conv2d(512, self.hyperInChn, 1, padding=(0, 0)),
-            nn.ReLU(inplace=True)
-        )
-
-        # Hyper network part, conv for generating target fc weights, fc for generating target fc biases
-        self.fc1w_conv = nn.Conv2d(self.hyperInChn, int(self.target_in_size * self.f1 / feature_size ** 2), 3,  padding=(1, 1))
-        self.fc1b_fc = nn.Linear(self.hyperInChn, self.f1)
-
-        self.fc2w_conv = nn.Conv2d(self.hyperInChn, int(self.f1 * self.f2 / feature_size ** 2), 3, padding=(1, 1))
-        self.fc2b_fc = nn.Linear(self.hyperInChn, self.f2)
-
-        self.fc3w_conv = nn.Conv2d(self.hyperInChn, int(self.f2 * self.f3 / feature_size ** 2), 3, padding=(1, 1))
-        self.fc3b_fc = nn.Linear(self.hyperInChn, self.f3)
-
-        self.fc4w_conv = nn.Conv2d(self.hyperInChn, int(self.f3 * self.f4 / feature_size ** 2), 3, padding=(1, 1))
-        self.fc4b_fc = nn.Linear(self.hyperInChn, self.f4)
-
-        self.fc5w_fc = nn.Linear(self.hyperInChn, self.f4)
-        self.fc5b_fc = nn.Linear(self.hyperInChn, 1)
-
-        # initialize
-        for i, m_name in enumerate(self._modules):
-            if i > 2:
-                nn.init.kaiming_normal_(self._modules[m_name].weight.data)
-
-    def forward(self, img):
-        feature_size = self.feature_size
-
-        res_out = self.res(img)
-
-        # input vector for target net
-        target_in_vec = res_out['target_in_vec'].reshape(-1, self.target_in_size, 1, 1)
-
-        # input features for hyper net
-        hyper_in_feat = self.conv1(res_out['hyper_in_feat']).reshape(-1, self.hyperInChn, feature_size, feature_size)
-
-        # generating target net weights & biases
-        target_fc1w = self.fc1w_conv(hyper_in_feat).reshape(-1, self.f1, self.target_in_size, 1, 1)
-        target_fc1b = self.fc1b_fc(self.pool(hyper_in_feat).squeeze()).reshape(-1, self.f1)
-
-        target_fc2w = self.fc2w_conv(hyper_in_feat).reshape(-1, self.f2, self.f1, 1, 1)
-        target_fc2b = self.fc2b_fc(self.pool(hyper_in_feat).squeeze()).reshape(-1, self.f2)
-
-        target_fc3w = self.fc3w_conv(hyper_in_feat).reshape(-1, self.f3, self.f2, 1, 1)
-        target_fc3b = self.fc3b_fc(self.pool(hyper_in_feat).squeeze()).reshape(-1, self.f3)
-
-        target_fc4w = self.fc4w_conv(hyper_in_feat).reshape(-1, self.f4, self.f3, 1, 1)
-        target_fc4b = self.fc4b_fc(self.pool(hyper_in_feat).squeeze()).reshape(-1, self.f4)
-
-        target_fc5w = self.fc5w_fc(self.pool(hyper_in_feat).squeeze()).reshape(-1, 1, self.f4, 1, 1)
-        target_fc5b = self.fc5b_fc(self.pool(hyper_in_feat).squeeze()).reshape(-1, 1)
-
-        out = {}
-        out['target_in_vec'] = target_in_vec
-        out['target_fc1w'] = target_fc1w
-        out['target_fc1b'] = target_fc1b
-        out['target_fc2w'] = target_fc2w
-        out['target_fc2b'] = target_fc2b
-        out['target_fc3w'] = target_fc3w
-        out['target_fc3b'] = target_fc3b
-        out['target_fc4w'] = target_fc4w
-        out['target_fc4b'] = target_fc4b
-        out['target_fc5w'] = target_fc5w
-        out['target_fc5b'] = target_fc5b
-
-        return out
-
-
-class TargetNet(nn.Module):
-    """
-    Target network for quality prediction.
-    """
-    def __init__(self, paras):
-        super(TargetNet, self).__init__()
-        self.l1 = nn.Sequential(
-            TargetFC(paras['target_fc1w'], paras['target_fc1b']),
-            nn.Sigmoid(),
-        )
-        self.l2 = nn.Sequential(
-            TargetFC(paras['target_fc2w'], paras['target_fc2b']),
-            nn.Sigmoid(),
-        )
-
-        self.l3 = nn.Sequential(
-            TargetFC(paras['target_fc3w'], paras['target_fc3b']),
-            nn.Sigmoid(),
-        )
-
-        self.l4 = nn.Sequential(
-            TargetFC(paras['target_fc4w'], paras['target_fc4b']),
-            nn.Sigmoid(),
-            TargetFC(paras['target_fc5w'], paras['target_fc5b']),
-        )
-
-    def forward(self, x):
-        q = self.l1(x)
-        # q = F.dropout(q)
-        q = self.l2(q)
-        q = self.l3(q)
-        q = self.l4(q).squeeze()
-        return q
-
-
-class TargetFC(nn.Module):
-    """
-    Fully connection operations for target net
-
-    Note:
-        Weights & biases are different for different images in a batch,
-        thus here we use group convolution for calculating images in a batch with individual weights & biases.
-    """
-    def __init__(self, weight, bias):
-        super(TargetFC, self).__init__()
-        self.weight = weight
-        self.bias = bias
-
-    def forward(self, input_):
-
-        input_re = input_.reshape(-1, input_.shape[0] * input_.shape[1], input_.shape[2], input_.shape[3])
-        weight_re = self.weight.reshape(self.weight.shape[0] * self.weight.shape[1], self.weight.shape[2], self.weight.shape[3], self.weight.shape[4])
-        bias_re = self.bias.reshape(self.bias.shape[0] * self.bias.shape[1])
-        out = F.conv2d(input=input_re, weight=weight_re, bias=bias_re, groups=self.weight.shape[0])
-
-        return out.reshape(input_.shape[0], self.weight.shape[1], input_.shape[2], input_.shape[3])
-
-
-class Bottleneck(nn.Module):
-    expansion = 4
-
-    def __init__(self, inplanes, planes, stride=1, downsample=None):
-        super(Bottleneck, self).__init__()
-        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
-        self.bn1 = nn.BatchNorm2d(planes)
-        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
-                               padding=1, bias=False)
-        self.bn2 = nn.BatchNorm2d(planes)
-        self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
-        self.bn3 = nn.BatchNorm2d(planes * 4)
-        self.relu = nn.ReLU(inplace=True)
-        self.downsample = downsample
-        self.stride = stride
-
-    def forward(self, x):
-        residual = x
-
-        out = self.conv1(x)
-        out = self.bn1(out)
-        out = self.relu(out)
-
-        out = self.conv2(out)
-        out = self.bn2(out)
-        out = self.relu(out)
-
-        out = self.conv3(out)
-        out = self.bn3(out)
-
-        if self.downsample is not None:
-            residual = self.downsample(x)
-
-        out += residual
-        out = self.relu(out)
-
-        return out
-
-
-class ResNetBackbone(nn.Module):
-
-    def __init__(self, lda_out_channels, in_chn, block, layers, num_classes=1000):
-        super(ResNetBackbone, self).__init__()
-        self.inplanes = 64
-        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
-        self.bn1 = nn.BatchNorm2d(64)
-        self.relu = nn.ReLU(inplace=True)
-        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
-        self.layer1 = self._make_layer(block, 64, layers[0])
-        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
-        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
-        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
-
-        # local distortion aware module
-        self.lda1_pool = nn.Sequential(
-            nn.Conv2d(256, 16, kernel_size=1, stride=1, padding=0, bias=False),
-            nn.AvgPool2d(7, stride=7),
-        )
-        self.lda1_fc = nn.Linear(16 * 64, lda_out_channels)
-
-        self.lda2_pool = nn.Sequential(
-            nn.Conv2d(512, 32, kernel_size=1, stride=1, padding=0, bias=False),
-            nn.AvgPool2d(7, stride=7),
-        )
-        self.lda2_fc = nn.Linear(32 * 16, lda_out_channels)
-
-        self.lda3_pool = nn.Sequential(
-            nn.Conv2d(1024, 64, kernel_size=1, stride=1, padding=0, bias=False),
-            nn.AvgPool2d(7, stride=7),
-        )
-        self.lda3_fc = nn.Linear(64 * 4, lda_out_channels)
-
-        self.lda4_pool = nn.AvgPool2d(7, stride=7)
-        self.lda4_fc = nn.Linear(2048, in_chn - lda_out_channels * 3)
-
-        for m in self.modules():
-            if isinstance(m, nn.Conv2d):
-                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
-                m.weight.data.normal_(0, math.sqrt(2. / n))
-            elif isinstance(m, nn.BatchNorm2d):
-                m.weight.data.fill_(1)
-                m.bias.data.zero_()
-
-        # initialize
-        nn.init.kaiming_normal_(self.lda1_pool._modules['0'].weight.data)
-        nn.init.kaiming_normal_(self.lda2_pool._modules['0'].weight.data)
-        nn.init.kaiming_normal_(self.lda3_pool._modules['0'].weight.data)
-        nn.init.kaiming_normal_(self.lda1_fc.weight.data)
-        nn.init.kaiming_normal_(self.lda2_fc.weight.data)
-        nn.init.kaiming_normal_(self.lda3_fc.weight.data)
-        nn.init.kaiming_normal_(self.lda4_fc.weight.data)
-
-    def _make_layer(self, block, planes, blocks, stride=1):
-        downsample = None
-        if stride != 1 or self.inplanes != planes * block.expansion:
-            downsample = nn.Sequential(
-                nn.Conv2d(self.inplanes, planes * block.expansion,
-                          kernel_size=1, stride=stride, bias=False),
-                nn.BatchNorm2d(planes * block.expansion),
-            )
-
-        layers = []
-        layers.append(block(self.inplanes, planes, stride, downsample))
-        self.inplanes = planes * block.expansion
-        for i in range(1, blocks):
-            layers.append(block(self.inplanes, planes))
-
-        return nn.Sequential(*layers)
-
-    def forward(self, x):
-        x = self.conv1(x)
-        x = self.bn1(x)
-        x = self.relu(x)
-        x = self.maxpool(x)
-        x = self.layer1(x)
-
-        # the same effect as lda operation in the paper, but save much more memory
-        lda_1 = self.lda1_fc(self.lda1_pool(x).reshape(x.size(0), -1))
-        x = self.layer2(x)
-        lda_2 = self.lda2_fc(self.lda2_pool(x).reshape(x.size(0), -1))
-        x = self.layer3(x)
-        lda_3 = self.lda3_fc(self.lda3_pool(x).reshape(x.size(0), -1))
-        x = self.layer4(x)
-        lda_4 = self.lda4_fc(self.lda4_pool(x).reshape(x.size(0), -1))
-
-        vec = torch.cat((lda_1, lda_2, lda_3, lda_4), 1)
-
-        out = {}
-        out['hyper_in_feat'] = x
-        out['target_in_vec'] = vec
-
-        return out
-
-
-def resnet50_backbone(lda_out_channels, in_chn, pretrained=False, **kwargs):
-    """Constructs a ResNet-50 model_hyper.
-
-    Args:
-        pretrained (bool): If True, returns a model_hyper pre-trained on ImageNet
-    """
-    model = ResNetBackbone(lda_out_channels, in_chn, Bottleneck, [3, 4, 6, 3], **kwargs)
-    if pretrained:
-        save_model = model_zoo.load_url(model_urls['resnet50'])
-        model_dict = model.state_dict()
-        state_dict = {k: v for k, v in save_model.items() if k in model_dict.keys()}
-        model_dict.update(state_dict)
-        model.load_state_dict(model_dict)
-    else:
-        model.apply(weights_init_xavier)
-    return model
-
-
-def weights_init_xavier(m):
-    classname = m.__class__.__name__
-    # print(classname)
-    # if isinstance(m, nn.Conv2d):
-    if classname.find('Conv') != -1:
-        init.kaiming_normal_(m.weight.data)
-    elif classname.find('Linear') != -1:
-        init.kaiming_normal_(m.weight.data)
-    elif classname.find('BatchNorm2d') != -1:
-        init.uniform_(m.weight.data, 1.0, 0.02)
-        init.constant_(m.bias.data, 0.0)

eval/inference_videos.py

@@ -1,77 +0,0 @@
-# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import os
-import subprocess
-from tqdm import tqdm
-import random
-
-
-def inference_video_from_fileslist(
-    video_fileslist: str,
-    audio_fileslist: str,
-    output_dir: str,
-    unet_config_path: str,
-    ckpt_path: str,
-    guidance_scale: float,
-    seed: int = 42,
-):
-    with open(video_fileslist, "r", encoding="utf-8") as file:
-        video_paths = [line.strip() for line in file.readlines()]
-
-    with open(audio_fileslist, "r", encoding="utf-8") as file:
-        audio_paths = [line.strip() for line in file.readlines()]
-
-    random.seed(seed)
-
-    output_dir = f"{output_dir}__{seed}"
-    os.makedirs(output_dir, exist_ok=True)
-
-    random.shuffle(video_paths)
-    random.shuffle(audio_paths)
-
-    min_length = min(len(video_paths), len(audio_paths))
-
-    video_paths = video_paths[:min_length]
-    audio_paths = audio_paths[:min_length]
-
-    random.shuffle(video_paths)
-    random.shuffle(audio_paths)
-
-    for index, video_path in tqdm(enumerate(video_paths), total=len(video_paths)):
-        audio_path = audio_paths[index]
-        video_name = os.path.basename(video_path)[:-4]
-        audio_name = os.path.basename(audio_path)[:-4]
-        video_out_path = os.path.join(output_dir, f"{video_name}__{audio_name}.mp4")
-        inference_command = (
-            f"python -m scripts.inference --enable_deepcache --guidance_scale {guidance_scale} --unet_config_path {unet_config_path} "
-            f"--video_path {video_path} --audio_path {audio_path} --video_out_path {video_out_path} --inference_ckpt_path {ckpt_path}"
-        )
-        subprocess.run(inference_command, shell=True)
-
-
-if __name__ == "__main__":
-    video_fileslist = "/mnt/bn/maliva-gen-ai-v2/chunyu.li/fileslist/video_fileslist.txt"
-    audio_fileslist = "/mnt/bn/maliva-gen-ai-v2/chunyu.li/fileslist/audio_fileslist.txt"
-    output_dir = "/mnt/bn/maliva-gen-ai-v2/chunyu.li/inference_videos_results"
-
-    unet_config_path = "configs/unet/stage2_512.yaml"
-    ckpt_path = "checkpoints/latentsync_unet.pt"
-    guidance_scale = 1.5
-
-    seed = 42
-
-    inference_video_from_fileslist(
-        video_fileslist, audio_fileslist, output_dir, unet_config_path, ckpt_path, guidance_scale, seed
-    )

eval/syncnet/__init__.py

@@ -1 +0,0 @@
-from .syncnet_eval import SyncNetEval

eval/syncnet/syncnet.py

@@ -1,113 +0,0 @@
-# https://github.com/joonson/syncnet_python/blob/master/SyncNetModel.py
-
-import torch
-import torch.nn as nn
-
-
-def save(model, filename):
-    with open(filename, "wb") as f:
-        torch.save(model, f)
-        print("%s saved." % filename)
-
-
-def load(filename):
-    net = torch.load(filename)
-    return net
-
-
-class S(nn.Module):
-    def __init__(self, num_layers_in_fc_layers=1024):
-        super(S, self).__init__()
-
-        self.__nFeatures__ = 24
-        self.__nChs__ = 32
-        self.__midChs__ = 32
-
-        self.netcnnaud = nn.Sequential(
-            nn.Conv2d(1, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
-            nn.BatchNorm2d(64),
-            nn.ReLU(inplace=True),
-            nn.MaxPool2d(kernel_size=(1, 1), stride=(1, 1)),
-            nn.Conv2d(64, 192, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
-            nn.BatchNorm2d(192),
-            nn.ReLU(inplace=True),
-            nn.MaxPool2d(kernel_size=(3, 3), stride=(1, 2)),
-            nn.Conv2d(192, 384, kernel_size=(3, 3), padding=(1, 1)),
-            nn.BatchNorm2d(384),
-            nn.ReLU(inplace=True),
-            nn.Conv2d(384, 256, kernel_size=(3, 3), padding=(1, 1)),
-            nn.BatchNorm2d(256),
-            nn.ReLU(inplace=True),
-            nn.Conv2d(256, 256, kernel_size=(3, 3), padding=(1, 1)),
-            nn.BatchNorm2d(256),
-            nn.ReLU(inplace=True),
-            nn.MaxPool2d(kernel_size=(3, 3), stride=(2, 2)),
-            nn.Conv2d(256, 512, kernel_size=(5, 4), padding=(0, 0)),
-            nn.BatchNorm2d(512),
-            nn.ReLU(),
-        )
-
-        self.netfcaud = nn.Sequential(
-            nn.Linear(512, 512),
-            nn.BatchNorm1d(512),
-            nn.ReLU(),
-            nn.Linear(512, num_layers_in_fc_layers),
-        )
-
-        self.netfclip = nn.Sequential(
-            nn.Linear(512, 512),
-            nn.BatchNorm1d(512),
-            nn.ReLU(),
-            nn.Linear(512, num_layers_in_fc_layers),
-        )
-
-        self.netcnnlip = nn.Sequential(
-            nn.Conv3d(3, 96, kernel_size=(5, 7, 7), stride=(1, 2, 2), padding=0),
-            nn.BatchNorm3d(96),
-            nn.ReLU(inplace=True),
-            nn.MaxPool3d(kernel_size=(1, 3, 3), stride=(1, 2, 2)),
-            nn.Conv3d(96, 256, kernel_size=(1, 5, 5), stride=(1, 2, 2), padding=(0, 1, 1)),
-            nn.BatchNorm3d(256),
-            nn.ReLU(inplace=True),
-            nn.MaxPool3d(kernel_size=(1, 3, 3), stride=(1, 2, 2), padding=(0, 1, 1)),
-            nn.Conv3d(256, 256, kernel_size=(1, 3, 3), padding=(0, 1, 1)),
-            nn.BatchNorm3d(256),
-            nn.ReLU(inplace=True),
-            nn.Conv3d(256, 256, kernel_size=(1, 3, 3), padding=(0, 1, 1)),
-            nn.BatchNorm3d(256),
-            nn.ReLU(inplace=True),
-            nn.Conv3d(256, 256, kernel_size=(1, 3, 3), padding=(0, 1, 1)),
-            nn.BatchNorm3d(256),
-            nn.ReLU(inplace=True),
-            nn.MaxPool3d(kernel_size=(1, 3, 3), stride=(1, 2, 2)),
-            nn.Conv3d(256, 512, kernel_size=(1, 6, 6), padding=0),
-            nn.BatchNorm3d(512),
-            nn.ReLU(inplace=True),
-        )
-
-    def forward_aud(self, x):
-
-        mid = self.netcnnaud(x)
-        # N x ch x 24 x M
-        mid = mid.view((mid.size()[0], -1))
-        # N x (ch x 24)
-        out = self.netfcaud(mid)
-
-        return out
-
-    def forward_lip(self, x):
-
-        mid = self.netcnnlip(x)
-        mid = mid.view((mid.size()[0], -1))
-        # N x (ch x 24)
-        out = self.netfclip(mid)
-
-        return out
-
-    def forward_lipfeat(self, x):
-
-        mid = self.netcnnlip(x)
-        out = mid.view((mid.size()[0], -1))
-        # N x (ch x 24)
-
-        return out

eval/syncnet/syncnet_eval.py

@@ -1,220 +0,0 @@
-# Adapted from https://github.com/joonson/syncnet_python/blob/master/SyncNetInstance.py
-
-import torch
-import numpy
-import time, pdb, argparse, subprocess, os, math, glob
-import cv2
-import python_speech_features
-
-from scipy import signal
-from scipy.io import wavfile
-from .syncnet import S
-from shutil import rmtree
-from latentsync.utils.util import check_model_and_download
-
-
-# ==================== Get OFFSET ====================
-
-# Video 25 FPS, Audio 16000HZ
-
-
-def calc_pdist(feat1, feat2, vshift=10):
-    win_size = vshift * 2 + 1
-
-    feat2p = torch.nn.functional.pad(feat2, (0, 0, vshift, vshift))
-
-    dists = []
-
-    for i in range(0, len(feat1)):
-
-        dists.append(
-            torch.nn.functional.pairwise_distance(feat1[[i], :].repeat(win_size, 1), feat2p[i : i + win_size, :])
-        )
-
-    return dists
-
-
-# ==================== MAIN DEF ====================
-
-
-class SyncNetEval(torch.nn.Module):
-    def __init__(self, dropout=0, num_layers_in_fc_layers=1024, device="cpu"):
-        super().__init__()
-
-        self.__S__ = S(num_layers_in_fc_layers=num_layers_in_fc_layers).to(device)
-        self.device = device
-
-    def evaluate(self, video_path, temp_dir="temp", batch_size=20, vshift=15):
-
-        self.__S__.eval()
-
-        # ========== ==========
-        # Convert files
-        # ========== ==========
-
-        if os.path.exists(temp_dir):
-            rmtree(temp_dir)
-
-        os.makedirs(temp_dir)
-
-        # temp_video_path = os.path.join(temp_dir, "temp.mp4")
-        # command = f"ffmpeg -loglevel error -nostdin -y -i {video_path} -vf scale='224:224' {temp_video_path}"
-        # subprocess.call(command, shell=True)
-
-        command = f"ffmpeg -loglevel error -nostdin -y -i {video_path} -f image2 {os.path.join(temp_dir, '%06d.jpg')}"
-        subprocess.call(command, shell=True, stdout=None)
-
-        command = f"ffmpeg -loglevel error -nostdin -y -i {video_path} -async 1 -ac 1 -vn -acodec pcm_s16le -ar 16000 {os.path.join(temp_dir, 'audio.wav')}"
-        subprocess.call(command, shell=True, stdout=None)
-
-        # ========== ==========
-        # Load video
-        # ========== ==========
-
-        images = []
-
-        flist = glob.glob(os.path.join(temp_dir, "*.jpg"))
-        flist.sort()
-
-        for fname in flist:
-            img_input = cv2.imread(fname)
-            img_input = cv2.resize(img_input, (224, 224))  # HARD CODED, CHANGE BEFORE RELEASE
-            images.append(img_input)
-
-        im = numpy.stack(images, axis=3)
-        im = numpy.expand_dims(im, axis=0)
-        im = numpy.transpose(im, (0, 3, 4, 1, 2))
-
-        imtv = torch.autograd.Variable(torch.from_numpy(im.astype(float)).float())
-
-        # ========== ==========
-        # Load audio
-        # ========== ==========
-
-        sample_rate, audio = wavfile.read(os.path.join(temp_dir, "audio.wav"))
-        mfcc = zip(*python_speech_features.mfcc(audio, sample_rate))
-        mfcc = numpy.stack([numpy.array(i) for i in mfcc])
-
-        cc = numpy.expand_dims(numpy.expand_dims(mfcc, axis=0), axis=0)
-        cct = torch.autograd.Variable(torch.from_numpy(cc.astype(float)).float())
-
-        # ========== ==========
-        # Check audio and video input length
-        # ========== ==========
-
-        # if (float(len(audio)) / 16000) != (float(len(images)) / 25):
-        #     print(
-        #         "WARNING: Audio (%.4fs) and video (%.4fs) lengths are different."
-        #         % (float(len(audio)) / 16000, float(len(images)) / 25)
-        #     )
-
-        min_length = min(len(images), math.floor(len(audio) / 640))
-
-        # ========== ==========
-        # Generate video and audio feats
-        # ========== ==========
-
-        lastframe = min_length - 5
-        im_feat = []
-        cc_feat = []
-
-        tS = time.time()
-        for i in range(0, lastframe, batch_size):
-
-            im_batch = [imtv[:, :, vframe : vframe + 5, :, :] for vframe in range(i, min(lastframe, i + batch_size))]
-            im_in = torch.cat(im_batch, 0)
-            im_out = self.__S__.forward_lip(im_in.to(self.device))
-            im_feat.append(im_out.data.cpu())
-
-            cc_batch = [
-                cct[:, :, :, vframe * 4 : vframe * 4 + 20] for vframe in range(i, min(lastframe, i + batch_size))
-            ]
-            cc_in = torch.cat(cc_batch, 0)
-            cc_out = self.__S__.forward_aud(cc_in.to(self.device))
-            cc_feat.append(cc_out.data.cpu())
-
-        im_feat = torch.cat(im_feat, 0)
-        cc_feat = torch.cat(cc_feat, 0)
-
-        # ========== ==========
-        # Compute offset
-        # ========== ==========
-
-        dists = calc_pdist(im_feat, cc_feat, vshift=vshift)
-        mean_dists = torch.mean(torch.stack(dists, 1), 1)
-
-        min_dist, minidx = torch.min(mean_dists, 0)
-
-        av_offset = vshift - minidx
-        conf = torch.median(mean_dists) - min_dist
-
-        fdist = numpy.stack([dist[minidx].numpy() for dist in dists])
-        # fdist   = numpy.pad(fdist, (3,3), 'constant', constant_values=15)
-        fconf = torch.median(mean_dists).numpy() - fdist
-        framewise_conf = signal.medfilt(fconf, kernel_size=9)
-
-        # numpy.set_printoptions(formatter={"float": "{: 0.3f}".format})
-        rmtree(temp_dir)
-        return av_offset.item(), min_dist.item(), conf.item()
-
-    def extract_feature(self, opt, videofile):
-
-        self.__S__.eval()
-
-        # ========== ==========
-        # Load video
-        # ========== ==========
-        cap = cv2.VideoCapture(videofile)
-
-        frame_num = 1
-        images = []
-        while frame_num:
-            frame_num += 1
-            ret, image = cap.read()
-            if ret == 0:
-                break
-
-            images.append(image)
-
-        im = numpy.stack(images, axis=3)
-        im = numpy.expand_dims(im, axis=0)
-        im = numpy.transpose(im, (0, 3, 4, 1, 2))
-
-        imtv = torch.autograd.Variable(torch.from_numpy(im.astype(float)).float())
-
-        # ========== ==========
-        # Generate video feats
-        # ========== ==========
-
-        lastframe = len(images) - 4
-        im_feat = []
-
-        tS = time.time()
-        for i in range(0, lastframe, opt.batch_size):
-
-            im_batch = [
-                imtv[:, :, vframe : vframe + 5, :, :] for vframe in range(i, min(lastframe, i + opt.batch_size))
-            ]
-            im_in = torch.cat(im_batch, 0)
-            im_out = self.__S__.forward_lipfeat(im_in.to(self.device))
-            im_feat.append(im_out.data.cpu())
-
-        im_feat = torch.cat(im_feat, 0)
-
-        # ========== ==========
-        # Compute offset
-        # ========== ==========
-
-        print("Compute time %.3f sec." % (time.time() - tS))
-
-        return im_feat
-
-    def loadParameters(self, path):
-        check_model_and_download(path)
-        loaded_state = torch.load(path, map_location=lambda storage, loc: storage, weights_only=True)
-
-        self_state = self.__S__.state_dict()
-
-        for name, param in loaded_state.items():
-
-            self_state[name].copy_(param)

eval/syncnet_detect.py

@@ -1,251 +0,0 @@
-# Adapted from https://github.com/joonson/syncnet_python/blob/master/run_pipeline.py
-
-import os, pdb, subprocess, glob, cv2
-import numpy as np
-from shutil import rmtree
-import torch
-
-from scenedetect.video_manager import VideoManager
-from scenedetect.scene_manager import SceneManager
-from scenedetect.stats_manager import StatsManager
-from scenedetect.detectors import ContentDetector
-
-from scipy.interpolate import interp1d
-from scipy.io import wavfile
-from scipy import signal
-
-from eval.detectors import S3FD
-
-
-class SyncNetDetector:
-    def __init__(self, device, detect_results_dir="detect_results"):
-        self.s3f_detector = S3FD(device=device)
-        self.detect_results_dir = detect_results_dir
-
-    def __call__(self, video_path: str, min_track=50, scale=False):
-        crop_dir = os.path.join(self.detect_results_dir, "crop")
-        video_dir = os.path.join(self.detect_results_dir, "video")
-        frames_dir = os.path.join(self.detect_results_dir, "frames")
-        temp_dir = os.path.join(self.detect_results_dir, "temp")
-
-        # ========== DELETE EXISTING DIRECTORIES ==========
-        if os.path.exists(crop_dir):
-            rmtree(crop_dir)
-
-        if os.path.exists(video_dir):
-            rmtree(video_dir)
-
-        if os.path.exists(frames_dir):
-            rmtree(frames_dir)
-
-        if os.path.exists(temp_dir):
-            rmtree(temp_dir)
-
-        # ========== MAKE NEW DIRECTORIES ==========
-
-        os.makedirs(crop_dir)
-        os.makedirs(video_dir)
-        os.makedirs(frames_dir)
-        os.makedirs(temp_dir)
-
-        # ========== CONVERT VIDEO AND EXTRACT FRAMES ==========
-
-        if scale:
-            scaled_video_path = os.path.join(video_dir, "scaled.mp4")
-            command = f"ffmpeg -loglevel error -y -nostdin -i {video_path} -vf scale='224:224' {scaled_video_path}"
-            subprocess.run(command, shell=True)
-            video_path = scaled_video_path
-
-        command = f"ffmpeg -y -nostdin -loglevel error -i {video_path} -qscale:v 2 -async 1 -r 25 {os.path.join(video_dir, 'video.mp4')}"
-        subprocess.run(command, shell=True, stdout=None)
-
-        command = f"ffmpeg -y -nostdin -loglevel error -i {os.path.join(video_dir, 'video.mp4')} -qscale:v 2 -f image2 {os.path.join(frames_dir, '%06d.jpg')}"
-        subprocess.run(command, shell=True, stdout=None)
-
-        command = f"ffmpeg -y -nostdin -loglevel error -i {os.path.join(video_dir, 'video.mp4')} -ac 1 -vn -acodec pcm_s16le -ar 16000 {os.path.join(video_dir, 'audio.wav')}"
-        subprocess.run(command, shell=True, stdout=None)
-
-        faces = self.detect_face(frames_dir)
-
-        scene = self.scene_detect(video_dir)
-
-        # Face tracking
-        alltracks = []
-
-        for shot in scene:
-            if shot[1].frame_num - shot[0].frame_num >= min_track:
-                alltracks.extend(self.track_face(faces[shot[0].frame_num : shot[1].frame_num], min_track=min_track))
-
-        # Face crop
-        for ii, track in enumerate(alltracks):
-            self.crop_video(track, os.path.join(crop_dir, "%05d" % ii), frames_dir, 25, temp_dir, video_dir)
-
-        rmtree(temp_dir)
-
-    def scene_detect(self, video_dir):
-        video_manager = VideoManager([os.path.join(video_dir, "video.mp4")])
-        stats_manager = StatsManager()
-        scene_manager = SceneManager(stats_manager)
-        # Add ContentDetector algorithm (constructor takes detector options like threshold).
-        scene_manager.add_detector(ContentDetector())
-        base_timecode = video_manager.get_base_timecode()
-
-        video_manager.set_downscale_factor()
-
-        video_manager.start()
-
-        scene_manager.detect_scenes(frame_source=video_manager)
-
-        scene_list = scene_manager.get_scene_list(base_timecode)
-
-        if scene_list == []:
-            scene_list = [(video_manager.get_base_timecode(), video_manager.get_current_timecode())]
-
-        return scene_list
-
-    def track_face(self, scenefaces, num_failed_det=25, min_track=50, min_face_size=100):
-
-        iouThres = 0.5  # Minimum IOU between consecutive face detections
-        tracks = []
-
-        while True:
-            track = []
-            for framefaces in scenefaces:
-                for face in framefaces:
-                    if track == []:
-                        track.append(face)
-                        framefaces.remove(face)
-                    elif face["frame"] - track[-1]["frame"] <= num_failed_det:
-                        iou = bounding_box_iou(face["bbox"], track[-1]["bbox"])
-                        if iou > iouThres:
-                            track.append(face)
-                            framefaces.remove(face)
-                            continue
-                    else:
-                        break
-
-            if track == []:
-                break
-            elif len(track) > min_track:
-
-                framenum = np.array([f["frame"] for f in track])
-                bboxes = np.array([np.array(f["bbox"]) for f in track])
-
-                frame_i = np.arange(framenum[0], framenum[-1] + 1)
-
-                bboxes_i = []
-                for ij in range(0, 4):
-                    interpfn = interp1d(framenum, bboxes[:, ij])
-                    bboxes_i.append(interpfn(frame_i))
-                bboxes_i = np.stack(bboxes_i, axis=1)
-
-                if (
-                    max(np.mean(bboxes_i[:, 2] - bboxes_i[:, 0]), np.mean(bboxes_i[:, 3] - bboxes_i[:, 1]))
-                    > min_face_size
-                ):
-                    tracks.append({"frame": frame_i, "bbox": bboxes_i})
-
-        return tracks
-
-    def detect_face(self, frames_dir, facedet_scale=0.25):
-        flist = glob.glob(os.path.join(frames_dir, "*.jpg"))
-        flist.sort()
-
-        dets = []
-
-        for fidx, fname in enumerate(flist):
-            image = cv2.imread(fname)
-
-            image_np = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
-            bboxes = self.s3f_detector.detect_faces(image_np, conf_th=0.9, scales=[facedet_scale])
-
-            dets.append([])
-            for bbox in bboxes:
-                dets[-1].append({"frame": fidx, "bbox": (bbox[:-1]).tolist(), "conf": bbox[-1]})
-
-        return dets
-
-    def crop_video(self, track, cropfile, frames_dir, frame_rate, temp_dir, video_dir, crop_scale=0.4):
-
-        flist = glob.glob(os.path.join(frames_dir, "*.jpg"))
-        flist.sort()
-
-        fourcc = cv2.VideoWriter_fourcc(*"mp4v")
-        vOut = cv2.VideoWriter(cropfile + "t.mp4", fourcc, frame_rate, (224, 224))
-
-        dets = {"x": [], "y": [], "s": []}
-
-        for det in track["bbox"]:
-
-            dets["s"].append(max((det[3] - det[1]), (det[2] - det[0])) / 2)
-            dets["y"].append((det[1] + det[3]) / 2)  # crop center x
-            dets["x"].append((det[0] + det[2]) / 2)  # crop center y
-
-        # Smooth detections
-        dets["s"] = signal.medfilt(dets["s"], kernel_size=13)
-        dets["x"] = signal.medfilt(dets["x"], kernel_size=13)
-        dets["y"] = signal.medfilt(dets["y"], kernel_size=13)
-
-        for fidx, frame in enumerate(track["frame"]):
-
-            cs = crop_scale
-
-            bs = dets["s"][fidx]  # Detection box size
-            bsi = int(bs * (1 + 2 * cs))  # Pad videos by this amount
-
-            image = cv2.imread(flist[frame])
-
-            frame = np.pad(image, ((bsi, bsi), (bsi, bsi), (0, 0)), "constant", constant_values=(110, 110))
-            my = dets["y"][fidx] + bsi  # BBox center Y
-            mx = dets["x"][fidx] + bsi  # BBox center X
-
-            face = frame[int(my - bs) : int(my + bs * (1 + 2 * cs)), int(mx - bs * (1 + cs)) : int(mx + bs * (1 + cs))]
-
-            vOut.write(cv2.resize(face, (224, 224)))
-
-        audiotmp = os.path.join(temp_dir, "audio.wav")
-        audiostart = (track["frame"][0]) / frame_rate
-        audioend = (track["frame"][-1] + 1) / frame_rate
-
-        vOut.release()
-
-        # ========== CROP AUDIO FILE ==========
-
-        command = "ffmpeg -y -nostdin -loglevel error -i %s -ss %.3f -to %.3f %s" % (
-            os.path.join(video_dir, "audio.wav"),
-            audiostart,
-            audioend,
-            audiotmp,
-        )
-        output = subprocess.run(command, shell=True, stdout=None)
-
-        sample_rate, audio = wavfile.read(audiotmp)
-
-        # ========== COMBINE AUDIO AND VIDEO FILES ==========
-
-        command = "ffmpeg -y -nostdin -loglevel error -i %st.mp4 -i %s -c:v copy -c:a aac %s.mp4" % (
-            cropfile,
-            audiotmp,
-            cropfile,
-        )
-        output = subprocess.run(command, shell=True, stdout=None)
-
-        os.remove(cropfile + "t.mp4")
-
-        return {"track": track, "proc_track": dets}
-
-
-def bounding_box_iou(boxA, boxB):
-    xA = max(boxA[0], boxB[0])
-    yA = max(boxA[1], boxB[1])
-    xB = min(boxA[2], boxB[2])
-    yB = min(boxA[3], boxB[3])
-
-    interArea = max(0, xB - xA) * max(0, yB - yA)
-
-    boxAArea = (boxA[2] - boxA[0]) * (boxA[3] - boxA[1])
-    boxBArea = (boxB[2] - boxB[0]) * (boxB[3] - boxB[1])
-
-    iou = interArea / float(boxAArea + boxBArea - interArea)
-
-    return iou

face_processing.py

@@ -1,585 +0,0 @@
-"""Face detection and region extraction for lipsync optimization (DEPRECATED - Pipeline handles this automatically)"""
-
-# NOTE: All functions in this module are DEPRECATED.
-# The lipsync pipeline (latentsync/pipelines/lipsync_pipeline.py) now handles:
-# - Face detection
-# - Affine transformation
-# - Crop
-# - Restore
-# These functions are kept for reference but not used in the new workflow.
-
-# import os
-# import math
-# import logging
-# from typing import List, Dict, Tuple, Optional
-#
-# import cv2
-# import numpy as np
-# import mediapipe as mp
-# from ffmpy import FFmpeg, FFRuntimeError
-#
-# from video_processing import get_video_info
-#
-# logger = logging.getLogger(__name__)
-#
-#
-# class FaceDetectionError(Exception):
-#     """Custom exception for face detection errors"""
-#
-#     pass
-#
-#
-# def sample_frames_from_video(
-#     video_path: str, output_dir: str, sample_count: int = 5
-# ) -> List[Tuple[int, str]]:
-#     """Extract uniform sample frames from video using OpenCV CUDA (HuggingFace)
-#
-#     Args:
-#         video_path: Path to video
-#         output_dir: Directory to save extracted frames
-#         sample_count: Number of frames to sample
-#
-#     Returns:
-#         List of (frame_index, frame_path) tuples
-#     """
-#     video_info = get_video_info(video_path)
-#     fps = video_info["fps"]
-#     duration = video_info["duration"]
-#     total_frames = int(duration * fps)
-#
-#     frames_dir = os.path.join(output_dir, "sampled_frames")
-#     os.makedirs(frames_dir, exist_ok=True)
-#
-#     if total_frames <= sample_count:
-#         frame_indices = list(range(total_frames))
-#     else:
-#         frame_indices = [
-#             int(i * total_frames / sample_count) for i in range(sample_count)
-#         ]
-#
-#     extracted_frames = []
-#     cap = cv2.VideoCapture(video_path)
-#
-#     try:
-#         for idx, frame_idx in enumerate(frame_indices):
-#             cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)
-#             ret, frame = cap.read()
-#
-#             if not ret or frame is None:
-#                 logger.warning(f"Failed to read frame {frame_idx}")
-#                 continue
-#
-#             frame_path = os.path.join(frames_dir, f"frame_{idx:04d}.jpg")
-#             cv2.imwrite(frame_path, frame, [cv2.IMWRITE_JPEG_QUALITY, 90])
-#             extracted_frames.append((frame_idx, frame_path))
-#     finally:
-#         cap.release()
-#
-#     logger.info(f"Extracted {len(extracted_frames)} frames from {video_path}")
-#     return extracted_frames
-#
-#
-# def detect_faces_in_frames(
-#     extracted_frames: List[Tuple[int, str]],
-#     min_confidence: float = 0.5,
-#     min_face_pixels: int = 100,
-# ) -> List[Dict]:
-#     """Detect faces in all sampled frames using MediaPipe Face Detection API
-#
-#     Args:
-#         extracted_frames: List of (frame_index, frame_path) tuples
-#         min_confidence: Minimum detection confidence (0-1)
-#         min_face_pixels: Minimum face size in pixels
-#
-#     Returns:
-#         List of detections: [{"frame_idx", "confidence", "bbox": (x, y, w, h)}]
-#     """
-#     detections = []
-#
-#     with mp.solutions.face_detection.FaceDetection(
-#         model_selection=0, min_detection_confidence=min_confidence
-#     ) as face_detection:
-#         for frame_idx, frame_path in extracted_frames:
-#             frame = cv2.imread(frame_path)
-#             if frame is None:
-#                 logger.warning(f"Failed to read frame: {frame_path}")
-#                 continue
-#
-#             h, w = frame.shape[:2]
-#             frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-#
-#             results = face_detection.process(frame_rgb)
-#
-#             if results.detections:
-#                 for detection in results.detections:
-#                     bbox = detection.location_data.relative_bounding_box
-#
-#                     x = int(bbox.xmin * w)
-#                     y = int(bbox.ymin * h)
-#                     face_w = int(bbox.width * w)
-#                     face_h = int(bbox.height * h)
-#
-#                     x = max(0, x)
-#                     y = max(0, y)
-#                     face_w = min(w - x, face_w)
-#                     face_h = min(h - y, face_h)
-#
-#                     confidence = detection.score[0] if detection.score else 0.0
-#
-#                     if face_w >= min_face_pixels and face_h >= min_face_pixels:
-#                         detections.append(
-#                             {
-#                                 "frame_idx": frame_idx,
-#                                 "confidence": float(confidence),
-#                                 "bbox": (x, y, face_w, face_h),
-#                             }
-#                         )
-#
-#     logger.info(f"Detected {len(detections)} faces in {len(extracted_frames)} frames")
-#     return detections
-#
-#
-# def cluster_face_detections(
-#     detections: List[Dict], max_distance: int = 100
-# ) -> List[List[Dict]]:
-#     """Group face detections belonging to the same person using clustering
-#
-#     Args:
-#         detections: List of face detections
-#         max_distance: Maximum distance (pixels) to consider detections as same person
-#
-#     Returns:
-#         List of clusters (each cluster is a list of detections)
-#     """
-#     if not detections:
-#         return []
-#
-#     clusters = []
-#     visited = set()
-#
-#     for i, det_i in enumerate(detections):
-#         if i in visited:
-#             continue
-#
-#         x_i, y_i, w_i, h_i = det_i["bbox"]
-#         center_i = (x_i + w_i / 2, y_i + h_i / 2)
-#
-#         cluster = [det_i]
-#         visited.add(i)
-#
-#         for j, det_j in enumerate(detections):
-#             if j in visited:
-#                 continue
-#
-#             x_j, y_j, w_j, h_j = det_j["bbox"]
-#             center_j = (x_j + w_j / 2, y_j + h_j / 2)
-#
-#             distance = math.sqrt(
-#                 (center_i[0] - center_j[0]) ** 2 + (center_i[1] - center_j[1]) ** 2
-#             )
-#
-#             if distance < max_distance:
-#                 cluster.append(det_j)
-#                 visited.add(j)
-#
-#         clusters.append(cluster)
-#
-#     logger.info(f"Clustered {len(detections)} detections into {len(clusters)} clusters")
-#     return clusters
-#
-#
-# def select_best_cluster(clusters: List[List[Dict]]) -> Optional[List[Dict]]:
-#     """Select the best face cluster (highest frequency)
-#
-#     Args:
-#         clusters: List of clusters
-#
-#     Returns:
-#         Best cluster (most frequent) or None
-#     """
-#     if not clusters:
-#         return None
-#
-#     scored_clusters = [(len(cluster), cluster) for cluster in clusters]
-#     scored_clusters.sort(key=lambda x: x[0], reverse=True)
-#
-#     best_cluster = scored_clusters[0][1]
-#     logger.info(f"Selected best cluster with {len(best_cluster)} detections")
-#     return best_cluster
-#
-#
-# def verify_face_stability(
-#     cluster: List[Dict], max_movement_percent: float = 0.3
-# ) -> bool:
-#     """Verify face doesn't move too much between frames
-#
-#     Args:
-#         cluster: Face detections for the same person
-#         max_movement_percent: Max movement as percentage of average face size
-#
-#     Returns:
-#         True if face is stable, False otherwise
-#     """
-#     if len(cluster) < 2:
-#         return True
-#
-#     centers = []
-#     sizes = []
-#
-#     for det in cluster:
-#         x, y, w, h = det["bbox"]
-#         centers.append((x + w / 2, y + h / 2))
-#         sizes.append(w * h)
-#
-#     avg_size = sum(sizes) / len(sizes)
-#     avg_face_dim = math.sqrt(avg_size)
-#     max_allowed_movement = avg_face_dim * max_movement_percent
-#
-#     for i in range(len(centers) - 1):
-#         dx = abs(centers[i + 1][0] - centers[i][0])
-#         dy = abs(centers[i + 1][1] - centers[i][1])
-#         movement = math.sqrt(dx**2 + dy**2)
-#
-#         if movement > max_allowed_movement:
-#             logger.warning(
-#                 f"Face movement {movement:.1f}px > {max_allowed_movement:.1f}px"
-#             )
-#             return False
-#
-#     return True
-#
-#
-# def calculate_face_bbox_from_cluster(cluster: List[Dict]) -> Dict:
-#     """Calculate average face bounding box from cluster
-#
-#     Args:
-#         cluster: Face detections for the same person
-#
-#     Returns:
-#         Dict: {"x", "y", "width", "height"}
-#     """
-#     weighted_x = 0
-#     weighted_y = 0
-#     weighted_w = 0
-#     weighted_h = 0
-#     total_weight = 0
-#
-#     for det in cluster:
-#         x, y, w, h = det["bbox"]
-#         weight = det["confidence"]
-#         weighted_x += x * weight
-#         weighted_y += y * weight
-#         weighted_w += w * weight
-#         weighted_h += h * weight
-#         total_weight += weight
-#
-#     avg_bbox = {
-#         "x": int(weighted_x / total_weight),
-#         "y": int(weighted_y / total_weight),
-#         "width": int(weighted_w / total_weight),
-#         "height": int(weighted_h / total_weight),
-#     }
-#
-#     return avg_bbox
-#
-#
-# def calculate_safe_crop_size(
-#     face_bbox: Dict, video_width: int, video_height: int, crop_size: int = 512
-# ) -> Dict:
-#     """Calculate safe crop region ensuring face is inside
-#
-#     Args:
-#         face_bbox: Face bounding box {"x", "y", "width", "height"}
-#         video_width: Video width
-#         video_height: Video height
-#         crop_size: Size of crop region (default: 512)
-#
-#     Returns:
-#         Dict: {"x", "y", "width", "height"}
-#     """
-#     crop_half = crop_size // 2
-#
-#     face_center_x = face_bbox["x"] + face_bbox["width"] / 2
-#     face_center_y = face_bbox["y"] + face_bbox["height"] / 2
-#
-#     crop_x = int(face_center_x - crop_half)
-#     crop_y = int(face_center_y - crop_half)
-#
-#     crop_x = max(0, crop_x)
-#     crop_y = max(0, crop_y)
-#     crop_x = min(video_width - crop_size, crop_x)
-#     crop_y = min(video_height - crop_size, crop_y)
-#
-#     face_right = face_bbox["x"] + face_bbox["width"]
-#     face_bottom = face_bbox["y"] + face_bbox["height"]
-#     crop_right = crop_x + crop_size
-#     crop_bottom = crop_y + crop_size
-#
-#     if (
-#         face_bbox["x"] < crop_x
-#         or face_bbox["y"] < crop_y
-#         or face_right > crop_right
-#         or face_bottom > crop_bottom
-#     ):
-#         if face_bbox["x"] < crop_x:
-#             crop_x = face_bbox["x"]
-#         elif face_right > crop_right:
-#             crop_x = face_right - crop_size
-#
-#         if face_bbox["y"] < crop_y:
-#             crop_y = face_bbox["y"]
-#         elif face_bottom > crop_bottom:
-#             crop_y = face_bottom - crop_size
-#
-#         crop_x = max(0, crop_x)
-#         crop_y = max(0, crop_y)
-#         crop_x = min(video_width - crop_size, crop_x)
-#         crop_y = min(video_height - crop_size, crop_y)
-#
-#     return {"x": crop_x, "y": crop_y, "width": crop_size, "height": crop_size}
-#
-#
-# def detect_face_region(
-#     video_path: str,
-#     output_dir: str,
-#     crop_size: int = 512,
-#     sample_count: int = 20,
-#     min_confidence: float = 0.5,
-#     min_face_pixels: int = 100,
-#     max_face_movement_percent: float = 0.3,
-# ) -> Dict:
-#     """Main function: Detect face and calculate safe crop (DEPRECATED - Pipeline handles this)
-#
-#     Args:
-#         video_path: Path to video
-#         output_dir: Directory for temporary files
-#         crop_size: Size of crop region (default: 512)
-#         sample_count: Number of frames to sample (default: 20)
-#         min_confidence: Minimum detection confidence
-#         min_face_pixels: Minimum face size in pixels
-#         max_face_movement_percent: Max allowed face movement
-#
-#     Returns:
-#         Dict: {"x", "y", "width", "height", "face_bbox"}
-#
-#     Raises:
-#         FaceDetectionError: If face detection fails
-#     """
-#     try:
-#         logger.info(f"Starting face detection for: {video_path}")
-#         video_info = get_video_info(video_path)
-#         video_w, video_h = video_info["width"], video_info["height"]
-#         logger.info(
-#             f"Video: {video_w}x{video_h}, {video_info['fps']:.1f}fps, {video_info['duration']:.1f}s"
-#         )
-#
-#         if video_w < crop_size or video_h < crop_size:
-#             raise FaceDetectionError(
-#                 f"Video resolution {video_w}x{video_h} < {crop_size}x{crop_size}. "
-#                 f"Please upload higher resolution video."
-#             )
-#
-#         extracted_frames = sample_frames_from_video(
-#             video_path, output_dir, sample_count
-#         )
-#         logger.info(f"Sampled {len(extracted_frames)} frames for detection")
-#
-#         detections = detect_faces_in_frames(
-#             extracted_frames, min_confidence, min_face_pixels
-#         )
-#
-#         if not detections:
-#             raise FaceDetectionError(
-#                 f"No face detected in {sample_count} sampled frames. "
-#                 f"Please upload a video with a visible face."
-#             )
-#
-#         logger.info(f"Found {len(detections)} face detections")
-#
-#         frames_with_face = len(set(d["frame_idx"] for d in detections))
-#         face_coverage = frames_with_face / len(extracted_frames)
-#         logger.info(
-#             f"Face coverage: {frames_with_face}/{len(extracted_frames)} ({face_coverage * 100:.1f}%)"
-#         )
-#
-#         if face_coverage < 0.5:
-#             raise FaceDetectionError(
-#                 f"Face detected in only {frames_with_face}/{len(extracted_frames)} frames "
-#                 f"({face_coverage * 100:.1f}%). "
-#                 f"Please upload a video with a visible face."
-#             )
-#
-#         clusters = cluster_face_detections(detections)
-#         logger.info(f"Grouped into {len(clusters)} face clusters")
-#
-#         best_cluster = select_best_cluster(clusters)
-#
-#         if best_cluster is None:
-#             raise FaceDetectionError(
-#                 f"Failed to identify main face in video. "
-#                 f"Please upload a video with a clear, visible face."
-#             )
-#
-#         logger.info(f"Selected main face cluster with {len(best_cluster)} detections")
-#
-#         if not verify_face_stability(best_cluster, max_face_movement_percent):
-#             raise FaceDetectionError(
-#                 f"Face moves too much between frames. "
-#                 f"Please upload a video with a stable face position."
-#             )
-#
-#         logger.info("Face stability check passed")
-#
-#         face_bbox = calculate_face_bbox_from_cluster(best_cluster)
-#         crop_bbox = calculate_safe_crop_size(face_bbox, video_w, video_h, crop_size)
-#
-#         crop_bbox["face_bbox"] = face_bbox
-#
-#         logger.info(
-#             f"Face detected at ({face_bbox['x']}, {face_bbox['y']}) "
-#             f"size {face_bbox['width']}x{face_bbox['height']}, "
-#             f"crop at ({crop_bbox['x']}, {crop_bbox['y']})"
-#         )
-#         logger.info("Face detection completed successfully")
-#
-#         return crop_bbox
-#
-#     except FaceDetectionError:
-#         raise
-#     except Exception as e:
-#         logger.error(f"Face detection failed: {e}")
-#         raise FaceDetectionError(f"Face detection failed: {str(e)}")
-#
-#
-# def crop_video_to_size(
-#     video_path: str, crop_bbox: Dict, output_dir: str, crop_size: int = 512
-# ) -> str:
-#     """Crop video to specified size using calculated bbox (DEPRECATED - Pipeline handles this)
-#
-#     Args:
-#         video_path: Path to input video
-#         crop_bbox: Crop region {"x", "y", "width", "height"}
-#         output_dir: Directory to save output
-#         crop_size: Size of crop region (default: 512)
-#
-#     Returns:
-#         Path to cropped video
-#     """
-#     output_path = os.path.join(output_dir, f"face_cropped_{crop_size}x{crop_size}.mp4")
-#
-#     logger.info(
-#         f"Crop box: x={crop_bbox['x']}, y={crop_bbox['y']}, "
-#         f"width={crop_bbox['width']}, height={crop_bbox['height']}"
-#     )
-#
-#     ffmpeg = FFmpeg(
-#         inputs={video_path: None},
-#         outputs={
-#             output_path: [
-#                 "-vf",
-#                 f"crop={crop_bbox['width']}:{crop_bbox['height']}:{crop_bbox['x']}:{crop_bbox['y']}",
-#                 "-c:v",
-#                 "libx264",
-#                 "-preset",
-#                 "slow",
-#                 "-crf",
-#                 "18",
-#                 "-profile:v",
-#                 "high",
-#                 "-pix_fmt",
-#                 "yuv420p",
-#                 "-c:a",
-#                 "copy",
-#                 "-loglevel",
-#                 "error",
-#                 "-y",
-#             ]
-#         },
-#     )
-#     try:
-#         ffmpeg.run()
-#     except FFRuntimeError as e:
-#         logger.error(f"FFmpeg failed: {e}")
-#         raise
-#     logger.info(f"Cropped video to {crop_size}x{crop_size}: {output_path}")
-#     return output_path
-#
-#
-# def blend_face_into_original(
-#     original_video: str,
-#     face_video: str,
-#     crop_bbox: Dict,
-#     output_dir: str,
-#     lipsynced_info: Dict | None = None,
-#     feather: int = 15,
-# ) -> str:
-#     """Blend face video back into original video with edge feather only (DEPRECATED - Pipeline handles this)
-#
-#     Args:
-#         original_video: Path to original video
-#         face_video: Path to lipsynced face video (cropped)
-#         crop_bbox: Crop region {"x", "y", "width", "height"}
-#         output_dir: Directory to save output
-#         lipsynced_info: Info of lipsynced video {"width", "height"} (optional)
-#         feather: Feather radius for smooth blending at edges
-#
-#     Returns:
-#         Path to blended video
-#     """
-#     output_path = os.path.join(output_dir, "face_blended.mp4")
-#
-#     overlay_x = crop_bbox["x"]
-#     overlay_y = crop_bbox["y"]
-#
-#     if lipsynced_info:
-#         face_width = lipsynced_info["width"]
-#         face_height = lipsynced_info["height"]
-#         logger.info(
-#             f"Blending {face_width}x{face_height} at ({overlay_x}, {overlay_y}) "
-#             f"(crop_bbox: {crop_bbox})"
-#         )
-#     else:
-#         face_width = crop_bbox["width"]
-#         face_height = crop_bbox["height"]
-#         logger.info(f"Blending at ({overlay_x}, {overlay_y})")
-#
-#     mask_w = face_width
-#     mask_h = face_height
-#
-#     feather_radius = 50
-#
-#     ffmpeg = FFmpeg(
-#         inputs={original_video: None, face_video: None},
-#         outputs={
-#             output_path: [
-#                 "-filter_complex",
-#                 f"[0:v][1:v]overlay={overlay_x}:{overlay_y}",
-#                 "-c:v",
-#                 "libx264",
-#                 "-preset",
-#                 "slow",
-#                 "-crf",
-#                 "18",
-#                 "-profile:v",
-#                 "high",
-#                 "-pix_fmt",
-#                 "yuv420p",
-#                 "-threads",
-#                 "0",
-#                 "-movflags",
-#                 "+faststart",
-#                 "-c:a",
-#                 "copy",
-#                 "-loglevel",
-#                 "error",
-#                 "-y",
-#             ]
-#         },
-#     )
-#     try:
-#         ffmpeg.run()
-#     except FFRuntimeError as e:
-#         logger.error(f"FFmpeg failed: {e}")
-#         raise
-#     logger.info(f"Blended face into original: {output_path}")
-#     return output_path

latentsync/data/syncnet_dataset.py

@@ -1,139 +0,0 @@
-# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import os
-import numpy as np
-from torch.utils.data import Dataset
-import torch
-import random
-from ..utils.util import gather_video_paths_recursively
-from ..utils.image_processor import ImageProcessor
-from ..utils.audio import melspectrogram
-import math
-from pathlib import Path
-
-from decord import AudioReader, VideoReader, cpu
-
-
-class SyncNetDataset(Dataset):
-    def __init__(self, data_dir: str, fileslist: str, config):
-        if fileslist != "":
-            with open(fileslist) as file:
-                self.video_paths = [line.rstrip() for line in file]
-        elif data_dir != "":
-            self.video_paths = gather_video_paths_recursively(data_dir)
-        else:
-            raise ValueError("data_dir and fileslist cannot be both empty")
-
-        self.resolution = config.data.resolution
-        self.num_frames = config.data.num_frames
-
-        self.mel_window_length = math.ceil(self.num_frames / 5 * 16)
-
-        self.audio_sample_rate = config.data.audio_sample_rate
-        self.video_fps = config.data.video_fps
-        self.image_processor = ImageProcessor(resolution=config.data.resolution)
-        self.audio_mel_cache_dir = config.data.audio_mel_cache_dir
-        Path(self.audio_mel_cache_dir).mkdir(parents=True, exist_ok=True)
-
-    def __len__(self):
-        return len(self.video_paths)
-
-    def read_audio(self, video_path: str):
-        ar = AudioReader(video_path, ctx=cpu(self.worker_id), sample_rate=self.audio_sample_rate)
-        original_mel = melspectrogram(ar[:].asnumpy().squeeze(0))
-        return torch.from_numpy(original_mel)
-
-    def crop_audio_window(self, original_mel, start_index):
-        start_idx = int(80.0 * (start_index / float(self.video_fps)))
-        end_idx = start_idx + self.mel_window_length
-        return original_mel[:, start_idx:end_idx].unsqueeze(0)
-
-    def get_frames(self, video_reader: VideoReader):
-        total_num_frames = len(video_reader)
-
-        start_idx = random.randint(0, total_num_frames - self.num_frames)
-        frames_index = np.arange(start_idx, start_idx + self.num_frames, dtype=int)
-
-        while True:
-            wrong_start_idx = random.randint(0, total_num_frames - self.num_frames)
-            if wrong_start_idx == start_idx:
-                continue
-            wrong_frames_index = np.arange(wrong_start_idx, wrong_start_idx + self.num_frames, dtype=int)
-            break
-
-        frames = video_reader.get_batch(frames_index).asnumpy()
-        wrong_frames = video_reader.get_batch(wrong_frames_index).asnumpy()
-
-        return frames, wrong_frames, start_idx
-
-    def worker_init_fn(self, worker_id):
-        self.worker_id = worker_id
-
-    def __getitem__(self, idx):
-        while True:
-            try:
-                idx = random.randint(0, len(self) - 1)
-
-                # Get video file path
-                video_path = self.video_paths[idx]
-
-                vr = VideoReader(video_path, ctx=cpu(self.worker_id))
-
-                if len(vr) < 2 * self.num_frames:
-                    continue
-
-                frames, wrong_frames, start_idx = self.get_frames(vr)
-
-                mel_cache_path = os.path.join(
-                    self.audio_mel_cache_dir, os.path.basename(video_path).replace(".mp4", "_mel.pt")
-                )
-
-                if os.path.isfile(mel_cache_path):
-                    try:
-                        original_mel = torch.load(mel_cache_path, weights_only=True)
-                    except Exception as e:
-                        print(f"{type(e).__name__} - {e} - {mel_cache_path}")
-                        os.remove(mel_cache_path)
-                        original_mel = self.read_audio(video_path)
-                        torch.save(original_mel, mel_cache_path)
-                else:
-                    original_mel = self.read_audio(video_path)
-                    torch.save(original_mel, mel_cache_path)
-
-                mel = self.crop_audio_window(original_mel, start_idx)
-
-                if mel.shape[-1] != self.mel_window_length:
-                    continue
-
-                if random.choice([True, False]):
-                    y = torch.ones(1).float()
-                    chosen_frames = frames
-                else:
-                    y = torch.zeros(1).float()
-                    chosen_frames = wrong_frames
-
-                chosen_frames = self.image_processor.process_images(chosen_frames)
-
-                vr.seek(0)  # avoid memory leak
-                break
-
-            except Exception as e:  # Handle the exception of face not detcted
-                print(f"{type(e).__name__} - {e} - {video_path}")
-                if "vr" in locals():
-                    vr.seek(0)  # avoid memory leak
-
-        sample = dict(frames=chosen_frames, audio_samples=mel, y=y)
-
-        return sample

latentsync/data/unet_dataset.py

@@ -1,152 +0,0 @@
-# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import os
-import math
-import numpy as np
-from torch.utils.data import Dataset
-import torch
-import random
-import cv2
-from ..utils.image_processor import ImageProcessor, load_fixed_mask
-from ..utils.audio import melspectrogram
-from decord import AudioReader, VideoReader, cpu
-import torch.nn.functional as F
-from pathlib import Path
-
-
-class UNetDataset(Dataset):
-    def __init__(self, train_data_dir: str, config):
-        if config.data.train_fileslist != "":
-            with open(config.data.train_fileslist) as file:
-                self.video_paths = [line.rstrip() for line in file]
-        elif train_data_dir != "":
-            self.video_paths = []
-            for file in os.listdir(train_data_dir):
-                if file.endswith(".mp4"):
-                    self.video_paths.append(os.path.join(train_data_dir, file))
-        else:
-            raise ValueError("data_dir and fileslist cannot be both empty")
-
-        self.resolution = config.data.resolution
-        self.num_frames = config.data.num_frames
-
-        self.mel_window_length = math.ceil(self.num_frames / 5 * 16)
-
-        self.audio_sample_rate = config.data.audio_sample_rate
-        self.video_fps = config.data.video_fps
-        self.image_processor = ImageProcessor(
-            self.resolution, mask_image=load_fixed_mask(self.resolution, config.data.mask_image_path)
-        )
-        self.load_audio_data = config.model.add_audio_layer and config.run.use_syncnet
-        self.audio_mel_cache_dir = config.data.audio_mel_cache_dir
-        Path(self.audio_mel_cache_dir).mkdir(parents=True, exist_ok=True)
-
-    def __len__(self):
-        return len(self.video_paths)
-
-    def read_audio(self, video_path: str):
-        ar = AudioReader(video_path, ctx=cpu(self.worker_id), sample_rate=self.audio_sample_rate)
-        original_mel = melspectrogram(ar[:].asnumpy().squeeze(0))
-        return torch.from_numpy(original_mel)
-
-    def crop_audio_window(self, original_mel, start_index):
-        start_idx = int(80.0 * (start_index / float(self.video_fps)))
-        end_idx = start_idx + self.mel_window_length
-        return original_mel[:, start_idx:end_idx].unsqueeze(0)
-
-    def get_frames(self, video_reader: VideoReader):
-        total_num_frames = len(video_reader)
-
-        start_idx = random.randint(0, total_num_frames - self.num_frames)
-        gt_frames_index = np.arange(start_idx, start_idx + self.num_frames, dtype=int)
-
-        while True:
-            ref_start_idx = random.randint(0, total_num_frames - self.num_frames)
-            if ref_start_idx > start_idx - self.num_frames and ref_start_idx < start_idx + self.num_frames:
-                continue
-            ref_frames_index = np.arange(ref_start_idx, ref_start_idx + self.num_frames, dtype=int)
-            break
-
-        gt_frames = video_reader.get_batch(gt_frames_index).asnumpy()
-        ref_frames = video_reader.get_batch(ref_frames_index).asnumpy()
-
-        return gt_frames, ref_frames, start_idx
-
-    def worker_init_fn(self, worker_id):
-        self.worker_id = worker_id
-
-    def __getitem__(self, idx):
-        while True:
-            try:
-                idx = random.randint(0, len(self) - 1)
-
-                # Get video file path
-                video_path = self.video_paths[idx]
-
-                vr = VideoReader(video_path, ctx=cpu(self.worker_id))
-
-                if len(vr) < 3 * self.num_frames:
-                    continue
-
-                gt_frames, ref_frames, start_idx = self.get_frames(vr)
-
-                if self.load_audio_data:
-                    mel_cache_path = os.path.join(
-                        self.audio_mel_cache_dir, os.path.basename(video_path).replace(".mp4", "_mel.pt")
-                    )
-
-                    if os.path.isfile(mel_cache_path):
-                        try:
-                            original_mel = torch.load(mel_cache_path, weights_only=True)
-                        except Exception as e:
-                            print(f"{type(e).__name__} - {e} - {mel_cache_path}")
-                            os.remove(mel_cache_path)
-                            original_mel = self.read_audio(video_path)
-                            torch.save(original_mel, mel_cache_path)
-                    else:
-                        original_mel = self.read_audio(video_path)
-                        torch.save(original_mel, mel_cache_path)
-
-                    mel = self.crop_audio_window(original_mel, start_idx)
-
-                    if mel.shape[-1] != self.mel_window_length:
-                        continue
-                else:
-                    mel = []
-
-                gt_pixel_values, masked_pixel_values, masks = self.image_processor.prepare_masks_and_masked_images(
-                    gt_frames, affine_transform=False
-                )  # (f, c, h, w)
-                ref_pixel_values = self.image_processor.process_images(ref_frames)
-
-                vr.seek(0)  # avoid memory leak
-                break
-
-            except Exception as e:  # Handle the exception of face not detcted
-                print(f"{type(e).__name__} - {e} - {video_path}")
-                if "vr" in locals():
-                    vr.seek(0)  # avoid memory leak
-
-        sample = dict(
-            gt_pixel_values=gt_pixel_values,
-            masked_pixel_values=masked_pixel_values,
-            ref_pixel_values=ref_pixel_values,
-            mel=mel,
-            masks=masks,
-            video_path=video_path,
-            start_idx=start_idx,
-        )
-
-        return sample

latentsync/models/attention.py

@@ -1,280 +0,0 @@
-# Adapted from https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention.py
-
-from dataclasses import dataclass
-from typing import Optional
-
-import torch
-import torch.nn.functional as F
-from torch import nn
-
-from diffusers.configuration_utils import ConfigMixin, register_to_config
-from diffusers.models import ModelMixin
-from diffusers.utils import BaseOutput
-from diffusers.models.attention import FeedForward, AdaLayerNorm
-
-from einops import rearrange, repeat
-
-
-@dataclass
-class Transformer3DModelOutput(BaseOutput):
-    sample: torch.FloatTensor
-
-
-class Transformer3DModel(ModelMixin, ConfigMixin):
-    @register_to_config
-    def __init__(
-        self,
-        num_attention_heads: int = 16,
-        attention_head_dim: int = 88,
-        in_channels: Optional[int] = None,
-        num_layers: int = 1,
-        dropout: float = 0.0,
-        norm_num_groups: int = 32,
-        cross_attention_dim: Optional[int] = None,
-        attention_bias: bool = False,
-        activation_fn: str = "geglu",
-        num_embeds_ada_norm: Optional[int] = None,
-        use_linear_projection: bool = False,
-        only_cross_attention: bool = False,
-        upcast_attention: bool = False,
-        add_audio_layer=False,
-    ):
-        super().__init__()
-        self.use_linear_projection = use_linear_projection
-        self.num_attention_heads = num_attention_heads
-        self.attention_head_dim = attention_head_dim
-        inner_dim = num_attention_heads * attention_head_dim
-
-        # Define input layers
-        self.in_channels = in_channels
-
-        self.norm = torch.nn.GroupNorm(num_groups=norm_num_groups, num_channels=in_channels, eps=1e-6, affine=True)
-        if use_linear_projection:
-            self.proj_in = nn.Linear(in_channels, inner_dim)
-        else:
-            self.proj_in = nn.Conv2d(in_channels, inner_dim, kernel_size=1, stride=1, padding=0)
-
-        # Define transformers blocks
-        self.transformer_blocks = nn.ModuleList(
-            [
-                BasicTransformerBlock(
-                    inner_dim,
-                    num_attention_heads,
-                    attention_head_dim,
-                    dropout=dropout,
-                    cross_attention_dim=cross_attention_dim,
-                    activation_fn=activation_fn,
-                    num_embeds_ada_norm=num_embeds_ada_norm,
-                    attention_bias=attention_bias,
-                    upcast_attention=upcast_attention,
-                    add_audio_layer=add_audio_layer,
-                )
-                for d in range(num_layers)
-            ]
-        )
-
-        # Define output layers
-        if use_linear_projection:
-            self.proj_out = nn.Linear(in_channels, inner_dim)
-        else:
-            self.proj_out = nn.Conv2d(inner_dim, in_channels, kernel_size=1, stride=1, padding=0)
-
-    def forward(self, hidden_states, encoder_hidden_states=None, timestep=None, return_dict: bool = True):
-        # Input
-        assert hidden_states.dim() == 5, f"Expected hidden_states to have ndim=5, but got ndim={hidden_states.dim()}."
-        video_length = hidden_states.shape[2]
-        hidden_states = rearrange(hidden_states, "b c f h w -> (b f) c h w")
-
-        batch, channel, height, weight = hidden_states.shape
-        residual = hidden_states
-
-        hidden_states = self.norm(hidden_states)
-        if not self.use_linear_projection:
-            hidden_states = self.proj_in(hidden_states)
-            inner_dim = hidden_states.shape[1]
-            hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch, height * weight, inner_dim)
-        else:
-            inner_dim = hidden_states.shape[1]
-            hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch, height * weight, inner_dim)
-            hidden_states = self.proj_in(hidden_states)
-
-        # Blocks
-        for block in self.transformer_blocks:
-            hidden_states = block(
-                hidden_states,
-                encoder_hidden_states=encoder_hidden_states,
-                timestep=timestep,
-                video_length=video_length,
-            )
-
-        # Output
-        if not self.use_linear_projection:
-            hidden_states = hidden_states.reshape(batch, height, weight, inner_dim).permute(0, 3, 1, 2).contiguous()
-            hidden_states = self.proj_out(hidden_states)
-        else:
-            hidden_states = self.proj_out(hidden_states)
-            hidden_states = hidden_states.reshape(batch, height, weight, inner_dim).permute(0, 3, 1, 2).contiguous()
-
-        output = hidden_states + residual
-
-        output = rearrange(output, "(b f) c h w -> b c f h w", f=video_length)
-        if not return_dict:
-            return (output,)
-
-        return Transformer3DModelOutput(sample=output)
-
-
-class BasicTransformerBlock(nn.Module):
-    def __init__(
-        self,
-        dim: int,
-        num_attention_heads: int,
-        attention_head_dim: int,
-        dropout=0.0,
-        cross_attention_dim: Optional[int] = None,
-        activation_fn: str = "geglu",
-        num_embeds_ada_norm: Optional[int] = None,
-        attention_bias: bool = False,
-        upcast_attention: bool = False,
-        add_audio_layer=False,
-    ):
-        super().__init__()
-        self.use_ada_layer_norm = num_embeds_ada_norm is not None
-        self.add_audio_layer = add_audio_layer
-
-        self.norm1 = AdaLayerNorm(dim, num_embeds_ada_norm) if self.use_ada_layer_norm else nn.LayerNorm(dim)
-        self.attn1 = Attention(
-            query_dim=dim,
-            heads=num_attention_heads,
-            dim_head=attention_head_dim,
-            dropout=dropout,
-            bias=attention_bias,
-            upcast_attention=upcast_attention,
-        )
-
-        # Cross-attn
-        if add_audio_layer:
-            self.norm2 = AdaLayerNorm(dim, num_embeds_ada_norm) if self.use_ada_layer_norm else nn.LayerNorm(dim)
-            self.attn2 = Attention(
-                query_dim=dim,
-                cross_attention_dim=cross_attention_dim,
-                heads=num_attention_heads,
-                dim_head=attention_head_dim,
-                dropout=dropout,
-                bias=attention_bias,
-                upcast_attention=upcast_attention,
-            )
-        else:
-            self.attn2 = None
-
-        # Feed-forward
-        self.ff = FeedForward(dim, dropout=dropout, activation_fn=activation_fn)
-        self.norm3 = nn.LayerNorm(dim)
-
-    def forward(
-        self, hidden_states, encoder_hidden_states=None, timestep=None, attention_mask=None, video_length=None
-    ):
-        norm_hidden_states = (
-            self.norm1(hidden_states, timestep) if self.use_ada_layer_norm else self.norm1(hidden_states)
-        )
-
-        hidden_states = self.attn1(norm_hidden_states, attention_mask=attention_mask) + hidden_states
-
-        if self.attn2 is not None and encoder_hidden_states is not None:
-            if encoder_hidden_states.dim() == 4:
-                encoder_hidden_states = rearrange(encoder_hidden_states, "b f s d -> (b f) s d")
-            norm_hidden_states = (
-                self.norm2(hidden_states, timestep) if self.use_ada_layer_norm else self.norm2(hidden_states)
-            )
-            hidden_states = (
-                self.attn2(
-                    norm_hidden_states, encoder_hidden_states=encoder_hidden_states, attention_mask=attention_mask
-                )
-                + hidden_states
-            )
-
-        # Feed-forward
-        hidden_states = self.ff(self.norm3(hidden_states)) + hidden_states
-
-        return hidden_states
-
-
-class Attention(nn.Module):
-    def __init__(
-        self,
-        query_dim: int,
-        cross_attention_dim: Optional[int] = None,
-        heads: int = 8,
-        dim_head: int = 64,
-        dropout: float = 0.0,
-        bias=False,
-        upcast_attention: bool = False,
-        upcast_softmax: bool = False,
-        norm_num_groups: Optional[int] = None,
-    ):
-        super().__init__()
-        inner_dim = dim_head * heads
-        cross_attention_dim = cross_attention_dim if cross_attention_dim is not None else query_dim
-        self.upcast_attention = upcast_attention
-        self.upcast_softmax = upcast_softmax
-
-        self.scale = dim_head**-0.5
-
-        self.heads = heads
-
-        if norm_num_groups is not None:
-            self.group_norm = nn.GroupNorm(num_channels=inner_dim, num_groups=norm_num_groups, eps=1e-5, affine=True)
-        else:
-            self.group_norm = None
-
-        self.to_q = nn.Linear(query_dim, inner_dim, bias=bias)
-        self.to_k = nn.Linear(cross_attention_dim, inner_dim, bias=bias)
-        self.to_v = nn.Linear(cross_attention_dim, inner_dim, bias=bias)
-
-        self.to_out = nn.ModuleList([])
-        self.to_out.append(nn.Linear(inner_dim, query_dim))
-        self.to_out.append(nn.Dropout(dropout))
-
-    def split_heads(self, tensor):
-        batch_size, seq_len, dim = tensor.shape
-        tensor = tensor.reshape(batch_size, seq_len, self.heads, dim // self.heads)
-        tensor = tensor.permute(0, 2, 1, 3)
-        return tensor
-
-    def concat_heads(self, tensor):
-        batch_size, heads, seq_len, head_dim = tensor.shape
-        tensor = tensor.permute(0, 2, 1, 3)
-        tensor = tensor.reshape(batch_size, seq_len, heads * head_dim)
-        return tensor
-
-    def forward(self, hidden_states, encoder_hidden_states=None, attention_mask=None):
-        if self.group_norm is not None:
-            hidden_states = self.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
-
-        query = self.to_q(hidden_states)
-        query = self.split_heads(query)
-
-        encoder_hidden_states = encoder_hidden_states if encoder_hidden_states is not None else hidden_states
-        key = self.to_k(encoder_hidden_states)
-        value = self.to_v(encoder_hidden_states)
-
-        key = self.split_heads(key)
-        value = self.split_heads(value)
-
-        if attention_mask is not None:
-            if attention_mask.shape[-1] != query.shape[1]:
-                target_length = query.shape[1]
-                attention_mask = F.pad(attention_mask, (0, target_length), value=0.0)
-                attention_mask = attention_mask.repeat_interleave(self.heads, dim=0)
-
-        # Use PyTorch native implementation of FlashAttention-2
-        hidden_states = F.scaled_dot_product_attention(query, key, value, attn_mask=attention_mask)
-
-        hidden_states = self.concat_heads(hidden_states)
-
-        # linear proj
-        hidden_states = self.to_out[0](hidden_states)
-
-        # dropout
-        hidden_states = self.to_out[1](hidden_states)
-        return hidden_states

latentsync/models/motion_module.py

@@ -1,313 +0,0 @@
-# Adapted from https://github.com/guoyww/AnimateDiff/blob/main/animatediff/models/motion_module.py
-
-# Actually we don't use the motion module in the final version of LatentSync
-# When we started the project, we used the codebase of AnimateDiff and tried motion module
-# But the results are poor, and we decied to leave the code here for possible future usage
-
-from dataclasses import dataclass
-
-import torch
-import torch.nn.functional as F
-from torch import nn
-
-from diffusers.configuration_utils import ConfigMixin, register_to_config
-from diffusers.models import ModelMixin
-from diffusers.utils import BaseOutput
-from diffusers.models.attention import FeedForward
-from .attention import Attention
-
-from einops import rearrange, repeat
-import math
-from .utils import zero_module
-
-
-@dataclass
-class TemporalTransformer3DModelOutput(BaseOutput):
-    sample: torch.FloatTensor
-
-
-def get_motion_module(in_channels, motion_module_type: str, motion_module_kwargs: dict):
-    if motion_module_type == "Vanilla":
-        return VanillaTemporalModule(
-            in_channels=in_channels,
-            **motion_module_kwargs,
-        )
-    else:
-        raise ValueError
-
-
-class VanillaTemporalModule(nn.Module):
-    def __init__(
-        self,
-        in_channels,
-        num_attention_heads=8,
-        num_transformer_block=2,
-        attention_block_types=("Temporal_Self", "Temporal_Self"),
-        cross_frame_attention_mode=None,
-        temporal_position_encoding=False,
-        temporal_position_encoding_max_len=24,
-        temporal_attention_dim_div=1,
-        zero_initialize=True,
-    ):
-        super().__init__()
-
-        self.temporal_transformer = TemporalTransformer3DModel(
-            in_channels=in_channels,
-            num_attention_heads=num_attention_heads,
-            attention_head_dim=in_channels // num_attention_heads // temporal_attention_dim_div,
-            num_layers=num_transformer_block,
-            attention_block_types=attention_block_types,
-            cross_frame_attention_mode=cross_frame_attention_mode,
-            temporal_position_encoding=temporal_position_encoding,
-            temporal_position_encoding_max_len=temporal_position_encoding_max_len,
-        )
-
-        if zero_initialize:
-            self.temporal_transformer.proj_out = zero_module(self.temporal_transformer.proj_out)
-
-    def forward(self, input_tensor, temb, encoder_hidden_states, attention_mask=None, anchor_frame_idx=None):
-        hidden_states = input_tensor
-        hidden_states = self.temporal_transformer(hidden_states, encoder_hidden_states, attention_mask)
-
-        output = hidden_states
-        return output
-
-
-class TemporalTransformer3DModel(nn.Module):
-    def __init__(
-        self,
-        in_channels,
-        num_attention_heads,
-        attention_head_dim,
-        num_layers,
-        attention_block_types=(
-            "Temporal_Self",
-            "Temporal_Self",
-        ),
-        dropout=0.0,
-        norm_num_groups=32,
-        cross_attention_dim=768,
-        activation_fn="geglu",
-        attention_bias=False,
-        upcast_attention=False,
-        cross_frame_attention_mode=None,
-        temporal_position_encoding=False,
-        temporal_position_encoding_max_len=24,
-    ):
-        super().__init__()
-
-        inner_dim = num_attention_heads * attention_head_dim
-
-        self.norm = torch.nn.GroupNorm(num_groups=norm_num_groups, num_channels=in_channels, eps=1e-6, affine=True)
-        self.proj_in = nn.Linear(in_channels, inner_dim)
-
-        self.transformer_blocks = nn.ModuleList(
-            [
-                TemporalTransformerBlock(
-                    dim=inner_dim,
-                    num_attention_heads=num_attention_heads,
-                    attention_head_dim=attention_head_dim,
-                    attention_block_types=attention_block_types,
-                    dropout=dropout,
-                    norm_num_groups=norm_num_groups,
-                    cross_attention_dim=cross_attention_dim,
-                    activation_fn=activation_fn,
-                    attention_bias=attention_bias,
-                    upcast_attention=upcast_attention,
-                    cross_frame_attention_mode=cross_frame_attention_mode,
-                    temporal_position_encoding=temporal_position_encoding,
-                    temporal_position_encoding_max_len=temporal_position_encoding_max_len,
-                )
-                for d in range(num_layers)
-            ]
-        )
-        self.proj_out = nn.Linear(inner_dim, in_channels)
-
-    def forward(self, hidden_states, encoder_hidden_states=None, attention_mask=None):
-        assert hidden_states.dim() == 5, f"Expected hidden_states to have ndim=5, but got ndim={hidden_states.dim()}."
-        video_length = hidden_states.shape[2]
-        hidden_states = rearrange(hidden_states, "b c f h w -> (b f) c h w")
-
-        batch, channel, height, weight = hidden_states.shape
-        residual = hidden_states
-
-        hidden_states = self.norm(hidden_states)
-        hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch, height * weight, channel)
-        hidden_states = self.proj_in(hidden_states)
-
-        # Transformer Blocks
-        for block in self.transformer_blocks:
-            hidden_states = block(
-                hidden_states, encoder_hidden_states=encoder_hidden_states, video_length=video_length
-            )
-
-        # output
-        hidden_states = self.proj_out(hidden_states)
-        hidden_states = hidden_states.reshape(batch, height, weight, channel).permute(0, 3, 1, 2).contiguous()
-
-        output = hidden_states + residual
-        output = rearrange(output, "(b f) c h w -> b c f h w", f=video_length)
-
-        return output
-
-
-class TemporalTransformerBlock(nn.Module):
-    def __init__(
-        self,
-        dim,
-        num_attention_heads,
-        attention_head_dim,
-        attention_block_types=(
-            "Temporal_Self",
-            "Temporal_Self",
-        ),
-        dropout=0.0,
-        norm_num_groups=32,
-        cross_attention_dim=768,
-        activation_fn="geglu",
-        attention_bias=False,
-        upcast_attention=False,
-        cross_frame_attention_mode=None,
-        temporal_position_encoding=False,
-        temporal_position_encoding_max_len=24,
-    ):
-        super().__init__()
-
-        attention_blocks = []
-        norms = []
-
-        for block_name in attention_block_types:
-            attention_blocks.append(
-                VersatileAttention(
-                    attention_mode=block_name.split("_")[0],
-                    cross_attention_dim=cross_attention_dim if block_name.endswith("_Cross") else None,
-                    query_dim=dim,
-                    heads=num_attention_heads,
-                    dim_head=attention_head_dim,
-                    dropout=dropout,
-                    bias=attention_bias,
-                    upcast_attention=upcast_attention,
-                    cross_frame_attention_mode=cross_frame_attention_mode,
-                    temporal_position_encoding=temporal_position_encoding,
-                    temporal_position_encoding_max_len=temporal_position_encoding_max_len,
-                )
-            )
-            norms.append(nn.LayerNorm(dim))
-
-        self.attention_blocks = nn.ModuleList(attention_blocks)
-        self.norms = nn.ModuleList(norms)
-
-        self.ff = FeedForward(dim, dropout=dropout, activation_fn=activation_fn)
-        self.ff_norm = nn.LayerNorm(dim)
-
-    def forward(self, hidden_states, encoder_hidden_states=None, attention_mask=None, video_length=None):
-        for attention_block, norm in zip(self.attention_blocks, self.norms):
-            norm_hidden_states = norm(hidden_states)
-            hidden_states = (
-                attention_block(
-                    norm_hidden_states,
-                    encoder_hidden_states=encoder_hidden_states if attention_block.is_cross_attention else None,
-                    video_length=video_length,
-                )
-                + hidden_states
-            )
-
-        hidden_states = self.ff(self.ff_norm(hidden_states)) + hidden_states
-
-        output = hidden_states
-        return output
-
-
-class PositionalEncoding(nn.Module):
-    def __init__(self, d_model, dropout=0.0, max_len=24):
-        super().__init__()
-        self.dropout = nn.Dropout(p=dropout)
-        position = torch.arange(max_len).unsqueeze(1)
-        div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))
-        pe = torch.zeros(1, max_len, d_model)
-        pe[0, :, 0::2] = torch.sin(position * div_term)
-        pe[0, :, 1::2] = torch.cos(position * div_term)
-        self.register_buffer("pe", pe)
-
-    def forward(self, x):
-        x = x + self.pe[:, : x.size(1)]
-        return self.dropout(x)
-
-
-class VersatileAttention(Attention):
-    def __init__(
-        self,
-        attention_mode=None,
-        cross_frame_attention_mode=None,
-        temporal_position_encoding=False,
-        temporal_position_encoding_max_len=24,
-        *args,
-        **kwargs,
-    ):
-        super().__init__(*args, **kwargs)
-        assert attention_mode == "Temporal"
-
-        self.attention_mode = attention_mode
-        self.is_cross_attention = kwargs["cross_attention_dim"] is not None
-
-        self.pos_encoder = (
-            PositionalEncoding(kwargs["query_dim"], dropout=0.0, max_len=temporal_position_encoding_max_len)
-            if (temporal_position_encoding and attention_mode == "Temporal")
-            else None
-        )
-
-    def extra_repr(self):
-        return f"(Module Info) Attention_Mode: {self.attention_mode}, Is_Cross_Attention: {self.is_cross_attention}"
-
-    def forward(self, hidden_states, encoder_hidden_states=None, attention_mask=None, video_length=None):
-        if self.attention_mode == "Temporal":
-            s = hidden_states.shape[1]
-            hidden_states = rearrange(hidden_states, "(b f) s c -> (b s) f c", f=video_length)
-
-            if self.pos_encoder is not None:
-                hidden_states = self.pos_encoder(hidden_states)
-
-            ##### This section will not be executed #####
-            encoder_hidden_states = (
-                repeat(encoder_hidden_states, "b n c -> (b s) n c", s=s)
-                if encoder_hidden_states is not None
-                else encoder_hidden_states
-            )
-            #############################################
-        else:
-            raise NotImplementedError
-
-        if self.group_norm is not None:
-            hidden_states = self.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
-
-        query = self.to_q(hidden_states)
-        query = self.split_heads(query)
-
-        encoder_hidden_states = encoder_hidden_states if encoder_hidden_states is not None else hidden_states
-        key = self.to_k(encoder_hidden_states)
-        value = self.to_v(encoder_hidden_states)
-
-        key = self.split_heads(key)
-        value = self.split_heads(value)
-
-        if attention_mask is not None:
-            if attention_mask.shape[-1] != query.shape[1]:
-                target_length = query.shape[1]
-                attention_mask = F.pad(attention_mask, (0, target_length), value=0.0)
-                attention_mask = attention_mask.repeat_interleave(self.heads, dim=0)
-
-        # Use PyTorch native implementation of FlashAttention-2
-        hidden_states = F.scaled_dot_product_attention(query, key, value, attn_mask=attention_mask)
-
-        hidden_states = self.concat_heads(hidden_states)
-
-        # linear proj
-        hidden_states = self.to_out[0](hidden_states)
-
-        # dropout
-        hidden_states = self.to_out[1](hidden_states)
-
-        if self.attention_mode == "Temporal":
-            hidden_states = rearrange(hidden_states, "(b s) f c -> (b f) s c", s=s)
-
-        return hidden_states

latentsync/models/resnet.py

@@ -1,228 +0,0 @@
-# Adapted from https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/resnet.py
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-from einops import rearrange
-
-
-class InflatedConv3d(nn.Conv2d):
-    def forward(self, x):
-        video_length = x.shape[2]
-
-        x = rearrange(x, "b c f h w -> (b f) c h w")
-        x = super().forward(x)
-        x = rearrange(x, "(b f) c h w -> b c f h w", f=video_length)
-
-        return x
-
-
-class InflatedGroupNorm(nn.GroupNorm):
-    def forward(self, x):
-        video_length = x.shape[2]
-
-        x = rearrange(x, "b c f h w -> (b f) c h w")
-        x = super().forward(x)
-        x = rearrange(x, "(b f) c h w -> b c f h w", f=video_length)
-
-        return x
-
-
-class Upsample3D(nn.Module):
-    def __init__(self, channels, use_conv=False, use_conv_transpose=False, out_channels=None, name="conv"):
-        super().__init__()
-        self.channels = channels
-        self.out_channels = out_channels or channels
-        self.use_conv = use_conv
-        self.use_conv_transpose = use_conv_transpose
-        self.name = name
-
-        conv = None
-        if use_conv_transpose:
-            raise NotImplementedError
-        elif use_conv:
-            self.conv = InflatedConv3d(self.channels, self.out_channels, 3, padding=1)
-
-    def forward(self, hidden_states, output_size=None):
-        assert hidden_states.shape[1] == self.channels
-
-        if self.use_conv_transpose:
-            raise NotImplementedError
-
-        # Cast to float32 to as 'upsample_nearest2d_out_frame' op does not support bfloat16
-        dtype = hidden_states.dtype
-        if dtype == torch.bfloat16:
-            hidden_states = hidden_states.to(torch.float32)
-
-        # upsample_nearest_nhwc fails with large batch sizes. see https://github.com/huggingface/diffusers/issues/984
-        if hidden_states.shape[0] >= 64:
-            hidden_states = hidden_states.contiguous()
-
-        # if `output_size` is passed we force the interpolation output
-        # size and do not make use of `scale_factor=2`
-        if output_size is None:
-            hidden_states = F.interpolate(hidden_states, scale_factor=[1.0, 2.0, 2.0], mode="nearest")
-        else:
-            hidden_states = F.interpolate(hidden_states, size=output_size, mode="nearest")
-
-        # If the input is bfloat16, we cast back to bfloat16
-        if dtype == torch.bfloat16:
-            hidden_states = hidden_states.to(dtype)
-
-        hidden_states = self.conv(hidden_states)
-
-        return hidden_states
-
-
-class Downsample3D(nn.Module):
-    def __init__(self, channels, use_conv=False, out_channels=None, padding=1, name="conv"):
-        super().__init__()
-        self.channels = channels
-        self.out_channels = out_channels or channels
-        self.use_conv = use_conv
-        self.padding = padding
-        stride = 2
-        self.name = name
-
-        if use_conv:
-            self.conv = InflatedConv3d(self.channels, self.out_channels, 3, stride=stride, padding=padding)
-        else:
-            raise NotImplementedError
-
-    def forward(self, hidden_states):
-        assert hidden_states.shape[1] == self.channels
-        if self.use_conv and self.padding == 0:
-            raise NotImplementedError
-
-        assert hidden_states.shape[1] == self.channels
-        hidden_states = self.conv(hidden_states)
-
-        return hidden_states
-
-
-class ResnetBlock3D(nn.Module):
-    def __init__(
-        self,
-        *,
-        in_channels,
-        out_channels=None,
-        conv_shortcut=False,
-        dropout=0.0,
-        temb_channels=512,
-        groups=32,
-        groups_out=None,
-        pre_norm=True,
-        eps=1e-6,
-        non_linearity="swish",
-        time_embedding_norm="default",
-        output_scale_factor=1.0,
-        use_in_shortcut=None,
-        use_inflated_groupnorm=False,
-    ):
-        super().__init__()
-        self.pre_norm = pre_norm
-        self.pre_norm = True
-        self.in_channels = in_channels
-        out_channels = in_channels if out_channels is None else out_channels
-        self.out_channels = out_channels
-        self.use_conv_shortcut = conv_shortcut
-        self.time_embedding_norm = time_embedding_norm
-        self.output_scale_factor = output_scale_factor
-
-        if groups_out is None:
-            groups_out = groups
-
-        assert use_inflated_groupnorm != None
-        if use_inflated_groupnorm:
-            self.norm1 = InflatedGroupNorm(num_groups=groups, num_channels=in_channels, eps=eps, affine=True)
-        else:
-            self.norm1 = torch.nn.GroupNorm(num_groups=groups, num_channels=in_channels, eps=eps, affine=True)
-
-        self.conv1 = InflatedConv3d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
-
-        if temb_channels is not None:
-            if self.time_embedding_norm == "default":
-                time_emb_proj_out_channels = out_channels
-            elif self.time_embedding_norm == "scale_shift":
-                time_emb_proj_out_channels = out_channels * 2
-            else:
-                raise ValueError(f"unknown time_embedding_norm : {self.time_embedding_norm} ")
-
-            self.time_emb_proj = torch.nn.Linear(temb_channels, time_emb_proj_out_channels)
-        else:
-            self.time_emb_proj = None
-
-        if self.time_embedding_norm == "scale_shift":
-            self.double_len_linear = torch.nn.Linear(time_emb_proj_out_channels, 2 * time_emb_proj_out_channels)
-        else:
-            self.double_len_linear = None
-
-        if use_inflated_groupnorm:
-            self.norm2 = InflatedGroupNorm(num_groups=groups_out, num_channels=out_channels, eps=eps, affine=True)
-        else:
-            self.norm2 = torch.nn.GroupNorm(num_groups=groups_out, num_channels=out_channels, eps=eps, affine=True)
-
-        self.dropout = torch.nn.Dropout(dropout)
-        self.conv2 = InflatedConv3d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
-
-        if non_linearity == "swish":
-            self.nonlinearity = lambda x: F.silu(x)
-        elif non_linearity == "mish":
-            self.nonlinearity = Mish()
-        elif non_linearity == "silu":
-            self.nonlinearity = nn.SiLU()
-
-        self.use_in_shortcut = self.in_channels != self.out_channels if use_in_shortcut is None else use_in_shortcut
-
-        self.conv_shortcut = None
-        if self.use_in_shortcut:
-            self.conv_shortcut = InflatedConv3d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
-
-    def forward(self, input_tensor, temb):
-        hidden_states = input_tensor
-
-        hidden_states = self.norm1(hidden_states)
-        hidden_states = self.nonlinearity(hidden_states)
-
-        hidden_states = self.conv1(hidden_states)
-
-        if temb is not None:
-            if temb.dim() == 2:
-                # input (1, 1280)
-                temb = self.time_emb_proj(self.nonlinearity(temb))
-                temb = temb[:, :, None, None, None]  # unsqueeze
-            else:
-                # input (1, 1280, 16)
-                temb = temb.permute(0, 2, 1)
-                temb = self.time_emb_proj(self.nonlinearity(temb))
-                if self.double_len_linear is not None:
-                    temb = self.double_len_linear(self.nonlinearity(temb))
-                temb = temb.permute(0, 2, 1)
-                temb = temb[:, :, :, None, None]
-
-        if temb is not None and self.time_embedding_norm == "default":
-            hidden_states = hidden_states + temb
-
-        hidden_states = self.norm2(hidden_states)
-
-        if temb is not None and self.time_embedding_norm == "scale_shift":
-            scale, shift = torch.chunk(temb, 2, dim=1)
-            hidden_states = hidden_states * (1 + scale) + shift
-
-        hidden_states = self.nonlinearity(hidden_states)
-
-        hidden_states = self.dropout(hidden_states)
-        hidden_states = self.conv2(hidden_states)
-
-        if self.conv_shortcut is not None:
-            input_tensor = self.conv_shortcut(input_tensor)
-
-        output_tensor = (input_tensor + hidden_states) / self.output_scale_factor
-
-        return output_tensor
-
-
-class Mish(torch.nn.Module):
-    def forward(self, hidden_states):
-        return hidden_states * torch.tanh(torch.nn.functional.softplus(hidden_states))

latentsync/models/stable_syncnet.py

@@ -1,233 +0,0 @@
-# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import torch
-from torch import nn
-from einops import rearrange
-from torch.nn import functional as F
-from .attention import Attention
-
-import torch.nn as nn
-import torch.nn.functional as F
-
-from diffusers.models.attention import FeedForward
-from einops import rearrange
-
-
-class StableSyncNet(nn.Module):
-    def __init__(self, config, gradient_checkpointing=False):
-        super().__init__()
-        self.audio_encoder = DownEncoder2D(
-            in_channels=config["audio_encoder"]["in_channels"],
-            block_out_channels=config["audio_encoder"]["block_out_channels"],
-            downsample_factors=config["audio_encoder"]["downsample_factors"],
-            dropout=config["audio_encoder"]["dropout"],
-            attn_blocks=config["audio_encoder"]["attn_blocks"],
-            gradient_checkpointing=gradient_checkpointing,
-        )
-
-        self.visual_encoder = DownEncoder2D(
-            in_channels=config["visual_encoder"]["in_channels"],
-            block_out_channels=config["visual_encoder"]["block_out_channels"],
-            downsample_factors=config["visual_encoder"]["downsample_factors"],
-            dropout=config["visual_encoder"]["dropout"],
-            attn_blocks=config["visual_encoder"]["attn_blocks"],
-            gradient_checkpointing=gradient_checkpointing,
-        )
-
-        self.eval()
-
-    def forward(self, image_sequences, audio_sequences):
-        vision_embeds = self.visual_encoder(image_sequences)  # (b, c, 1, 1)
-        audio_embeds = self.audio_encoder(audio_sequences)  # (b, c, 1, 1)
-
-        vision_embeds = vision_embeds.reshape(vision_embeds.shape[0], -1)  # (b, c)
-        audio_embeds = audio_embeds.reshape(audio_embeds.shape[0], -1)  # (b, c)
-
-        # Make them unit vectors
-        vision_embeds = F.normalize(vision_embeds, p=2, dim=1)
-        audio_embeds = F.normalize(audio_embeds, p=2, dim=1)
-
-        return vision_embeds, audio_embeds
-
-
-class ResnetBlock2D(nn.Module):
-    def __init__(
-        self,
-        in_channels: int,
-        out_channels: int,
-        dropout: float = 0.0,
-        norm_num_groups: int = 32,
-        eps: float = 1e-6,
-        act_fn: str = "silu",
-        downsample_factor=2,
-    ):
-        super().__init__()
-
-        self.norm1 = nn.GroupNorm(num_groups=norm_num_groups, num_channels=in_channels, eps=eps, affine=True)
-        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
-
-        self.norm2 = nn.GroupNorm(num_groups=norm_num_groups, num_channels=out_channels, eps=eps, affine=True)
-        self.dropout = nn.Dropout(dropout)
-        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
-
-        if act_fn == "relu":
-            self.act_fn = nn.ReLU()
-        elif act_fn == "silu":
-            self.act_fn = nn.SiLU()
-
-        if in_channels != out_channels:
-            self.conv_shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
-        else:
-            self.conv_shortcut = None
-
-        if isinstance(downsample_factor, list):
-            downsample_factor = tuple(downsample_factor)
-
-        if downsample_factor == 1:
-            self.downsample_conv = None
-        else:
-            self.downsample_conv = nn.Conv2d(
-                out_channels, out_channels, kernel_size=3, stride=downsample_factor, padding=0
-            )
-            self.pad = (0, 1, 0, 1)
-            if isinstance(downsample_factor, tuple):
-                if downsample_factor[0] == 1:
-                    self.pad = (0, 1, 1, 1)  # The padding order is from back to front
-                elif downsample_factor[1] == 1:
-                    self.pad = (1, 1, 0, 1)
-
-    def forward(self, input_tensor):
-        hidden_states = input_tensor
-
-        hidden_states = self.norm1(hidden_states)
-        hidden_states = self.act_fn(hidden_states)
-
-        hidden_states = self.conv1(hidden_states)
-        hidden_states = self.norm2(hidden_states)
-        hidden_states = self.act_fn(hidden_states)
-
-        hidden_states = self.dropout(hidden_states)
-        hidden_states = self.conv2(hidden_states)
-
-        if self.conv_shortcut is not None:
-            input_tensor = self.conv_shortcut(input_tensor)
-
-        hidden_states += input_tensor
-
-        if self.downsample_conv is not None:
-            hidden_states = F.pad(hidden_states, self.pad, mode="constant", value=0)
-            hidden_states = self.downsample_conv(hidden_states)
-
-        return hidden_states
-
-
-class AttentionBlock2D(nn.Module):
-    def __init__(self, query_dim, norm_num_groups=32, dropout=0.0):
-        super().__init__()
-        self.norm1 = torch.nn.GroupNorm(num_groups=norm_num_groups, num_channels=query_dim, eps=1e-6, affine=True)
-        self.norm2 = nn.LayerNorm(query_dim)
-        self.norm3 = nn.LayerNorm(query_dim)
-
-        self.ff = FeedForward(query_dim, dropout=dropout, activation_fn="geglu")
-
-        self.conv_in = nn.Conv2d(query_dim, query_dim, kernel_size=1, stride=1, padding=0)
-        self.conv_out = nn.Conv2d(query_dim, query_dim, kernel_size=1, stride=1, padding=0)
-
-        self.attn = Attention(query_dim=query_dim, heads=8, dim_head=query_dim // 8, dropout=dropout, bias=True)
-
-    def forward(self, hidden_states):
-        assert hidden_states.dim() == 4, f"Expected hidden_states to have ndim=4, but got ndim={hidden_states.dim()}."
-
-        batch, channel, height, width = hidden_states.shape
-        residual = hidden_states
-
-        hidden_states = self.norm1(hidden_states)
-        hidden_states = self.conv_in(hidden_states)
-        hidden_states = rearrange(hidden_states, "b c h w -> b (h w) c")
-
-        norm_hidden_states = self.norm2(hidden_states)
-
-        hidden_states = self.attn(norm_hidden_states, attention_mask=None) + hidden_states
-        hidden_states = self.ff(self.norm3(hidden_states)) + hidden_states
-
-        hidden_states = rearrange(hidden_states, "b (h w) c -> b c h w", h=height, w=width).contiguous()
-        hidden_states = self.conv_out(hidden_states)
-
-        hidden_states = hidden_states + residual
-        return hidden_states
-
-
-class DownEncoder2D(nn.Module):
-    def __init__(
-        self,
-        in_channels=4 * 16,
-        block_out_channels=[64, 128, 256, 256],
-        downsample_factors=[2, 2, 2, 2],
-        layers_per_block=2,
-        norm_num_groups=32,
-        attn_blocks=[1, 1, 1, 1],
-        dropout: float = 0.0,
-        act_fn="silu",
-        gradient_checkpointing=False,
-    ):
-        super().__init__()
-        self.layers_per_block = layers_per_block
-        self.gradient_checkpointing = gradient_checkpointing
-
-        # in
-        self.conv_in = nn.Conv2d(in_channels, block_out_channels[0], kernel_size=3, stride=1, padding=1)
-
-        # down
-        self.down_blocks = nn.ModuleList([])
-
-        output_channels = block_out_channels[0]
-        for i, block_out_channel in enumerate(block_out_channels):
-            input_channels = output_channels
-            output_channels = block_out_channel
-
-            down_block = ResnetBlock2D(
-                in_channels=input_channels,
-                out_channels=output_channels,
-                downsample_factor=downsample_factors[i],
-                norm_num_groups=norm_num_groups,
-                dropout=dropout,
-                act_fn=act_fn,
-            )
-
-            self.down_blocks.append(down_block)
-
-            if attn_blocks[i] == 1:
-                attention_block = AttentionBlock2D(query_dim=output_channels, dropout=dropout)
-                self.down_blocks.append(attention_block)
-
-        # out
-        self.norm_out = nn.GroupNorm(num_channels=block_out_channels[-1], num_groups=norm_num_groups, eps=1e-6)
-        self.act_fn_out = nn.ReLU()
-
-    def forward(self, hidden_states):
-        hidden_states = self.conv_in(hidden_states)
-
-        # down
-        for down_block in self.down_blocks:
-            if self.gradient_checkpointing:
-                hidden_states = torch.utils.checkpoint.checkpoint(down_block, hidden_states, use_reentrant=False)
-            else:
-                hidden_states = down_block(hidden_states)
-
-        # post-process
-        hidden_states = self.norm_out(hidden_states)
-        hidden_states = self.act_fn_out(hidden_states)
-
-        return hidden_states

latentsync/models/unet.py

@@ -1,512 +0,0 @@
-# Adapted from https://github.com/guoyww/AnimateDiff/blob/main/animatediff/models/unet.py
-
-from dataclasses import dataclass
-from typing import List, Optional, Tuple, Union
-import copy
-
-import torch
-import torch.nn as nn
-import torch.utils.checkpoint
-
-from diffusers.configuration_utils import ConfigMixin, register_to_config
-from diffusers.models import ModelMixin
-
-from diffusers.utils import BaseOutput, logging
-from diffusers.models.embeddings import TimestepEmbedding, Timesteps
-from .unet_blocks import (
-    CrossAttnDownBlock3D,
-    CrossAttnUpBlock3D,
-    DownBlock3D,
-    UNetMidBlock3DCrossAttn,
-    UpBlock3D,
-    get_down_block,
-    get_up_block,
-)
-from .resnet import InflatedConv3d, InflatedGroupNorm
-
-from ..utils.util import zero_rank_log
-from .utils import zero_module
-
-
-logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
-
-
-@dataclass
-class UNet3DConditionOutput(BaseOutput):
-    sample: torch.FloatTensor
-
-
-class UNet3DConditionModel(ModelMixin, ConfigMixin):
-    _supports_gradient_checkpointing = True
-
-    @register_to_config
-    def __init__(
-        self,
-        sample_size: Optional[int] = None,
-        in_channels: int = 4,
-        out_channels: int = 4,
-        center_input_sample: bool = False,
-        flip_sin_to_cos: bool = True,
-        freq_shift: int = 0,
-        down_block_types: Tuple[str] = (
-            "CrossAttnDownBlock3D",
-            "CrossAttnDownBlock3D",
-            "CrossAttnDownBlock3D",
-            "DownBlock3D",
-        ),
-        mid_block_type: str = "UNetMidBlock3DCrossAttn",
-        up_block_types: Tuple[str] = ("UpBlock3D", "CrossAttnUpBlock3D", "CrossAttnUpBlock3D", "CrossAttnUpBlock3D"),
-        only_cross_attention: Union[bool, Tuple[bool]] = False,
-        block_out_channels: Tuple[int] = (320, 640, 1280, 1280),
-        layers_per_block: int = 2,
-        downsample_padding: int = 1,
-        mid_block_scale_factor: float = 1,
-        act_fn: str = "silu",
-        norm_num_groups: int = 32,
-        norm_eps: float = 1e-5,
-        cross_attention_dim: int = 1280,
-        attention_head_dim: Union[int, Tuple[int]] = 8,
-        dual_cross_attention: bool = False,
-        use_linear_projection: bool = False,
-        class_embed_type: Optional[str] = None,
-        num_class_embeds: Optional[int] = None,
-        upcast_attention: bool = False,
-        resnet_time_scale_shift: str = "default",
-        use_inflated_groupnorm=False,
-        # Additional
-        use_motion_module=False,
-        motion_module_resolutions=(1, 2, 4, 8),
-        motion_module_mid_block=False,
-        motion_module_decoder_only=False,
-        motion_module_type=None,
-        motion_module_kwargs={},
-        add_audio_layer=False,
-    ):
-        super().__init__()
-
-        self.sample_size = sample_size
-        time_embed_dim = block_out_channels[0] * 4
-        self.use_motion_module = use_motion_module
-        self.add_audio_layer = add_audio_layer
-
-        self.conv_in = zero_module(InflatedConv3d(in_channels, block_out_channels[0], kernel_size=3, padding=(1, 1)))
-
-        # time
-        self.time_proj = Timesteps(block_out_channels[0], flip_sin_to_cos, freq_shift)
-        timestep_input_dim = block_out_channels[0]
-
-        self.time_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim)
-
-        # class embedding
-        if class_embed_type is None and num_class_embeds is not None:
-            self.class_embedding = nn.Embedding(num_class_embeds, time_embed_dim)
-        elif class_embed_type == "timestep":
-            self.class_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim)
-        elif class_embed_type == "identity":
-            self.class_embedding = nn.Identity(time_embed_dim, time_embed_dim)
-        else:
-            self.class_embedding = None
-
-        self.down_blocks = nn.ModuleList([])
-        self.mid_block = None
-        self.up_blocks = nn.ModuleList([])
-
-        if isinstance(only_cross_attention, bool):
-            only_cross_attention = [only_cross_attention] * len(down_block_types)
-
-        if isinstance(attention_head_dim, int):
-            attention_head_dim = (attention_head_dim,) * len(down_block_types)
-
-        # down
-        output_channel = block_out_channels[0]
-        for i, down_block_type in enumerate(down_block_types):
-            res = 2**i
-            input_channel = output_channel
-            output_channel = block_out_channels[i]
-            is_final_block = i == len(block_out_channels) - 1
-
-            down_block = get_down_block(
-                down_block_type,
-                num_layers=layers_per_block,
-                in_channels=input_channel,
-                out_channels=output_channel,
-                temb_channels=time_embed_dim,
-                add_downsample=not is_final_block,
-                resnet_eps=norm_eps,
-                resnet_act_fn=act_fn,
-                resnet_groups=norm_num_groups,
-                cross_attention_dim=cross_attention_dim,
-                attn_num_head_channels=attention_head_dim[i],
-                downsample_padding=downsample_padding,
-                dual_cross_attention=dual_cross_attention,
-                use_linear_projection=use_linear_projection,
-                only_cross_attention=only_cross_attention[i],
-                upcast_attention=upcast_attention,
-                resnet_time_scale_shift=resnet_time_scale_shift,
-                use_inflated_groupnorm=use_inflated_groupnorm,
-                use_motion_module=use_motion_module
-                and (res in motion_module_resolutions)
-                and (not motion_module_decoder_only),
-                motion_module_type=motion_module_type,
-                motion_module_kwargs=motion_module_kwargs,
-                add_audio_layer=add_audio_layer,
-            )
-            self.down_blocks.append(down_block)
-
-        # mid
-        if mid_block_type == "UNetMidBlock3DCrossAttn":
-            self.mid_block = UNetMidBlock3DCrossAttn(
-                in_channels=block_out_channels[-1],
-                temb_channels=time_embed_dim,
-                resnet_eps=norm_eps,
-                resnet_act_fn=act_fn,
-                output_scale_factor=mid_block_scale_factor,
-                resnet_time_scale_shift=resnet_time_scale_shift,
-                cross_attention_dim=cross_attention_dim,
-                attn_num_head_channels=attention_head_dim[-1],
-                resnet_groups=norm_num_groups,
-                dual_cross_attention=dual_cross_attention,
-                use_linear_projection=use_linear_projection,
-                upcast_attention=upcast_attention,
-                use_inflated_groupnorm=use_inflated_groupnorm,
-                use_motion_module=use_motion_module and motion_module_mid_block,
-                motion_module_type=motion_module_type,
-                motion_module_kwargs=motion_module_kwargs,
-                add_audio_layer=add_audio_layer,
-            )
-        else:
-            raise ValueError(f"unknown mid_block_type : {mid_block_type}")
-
-        # count how many layers upsample the videos
-        self.num_upsamplers = 0
-
-        # up
-        reversed_block_out_channels = list(reversed(block_out_channels))
-        reversed_attention_head_dim = list(reversed(attention_head_dim))
-        only_cross_attention = list(reversed(only_cross_attention))
-        output_channel = reversed_block_out_channels[0]
-        for i, up_block_type in enumerate(up_block_types):
-            res = 2 ** (3 - i)
-            is_final_block = i == len(block_out_channels) - 1
-
-            prev_output_channel = output_channel
-            output_channel = reversed_block_out_channels[i]
-            input_channel = reversed_block_out_channels[min(i + 1, len(block_out_channels) - 1)]
-
-            # add upsample block for all BUT final layer
-            if not is_final_block:
-                add_upsample = True
-                self.num_upsamplers += 1
-            else:
-                add_upsample = False
-
-            up_block = get_up_block(
-                up_block_type,
-                num_layers=layers_per_block + 1,
-                in_channels=input_channel,
-                out_channels=output_channel,
-                prev_output_channel=prev_output_channel,
-                temb_channels=time_embed_dim,
-                add_upsample=add_upsample,
-                resnet_eps=norm_eps,
-                resnet_act_fn=act_fn,
-                resnet_groups=norm_num_groups,
-                cross_attention_dim=cross_attention_dim,
-                attn_num_head_channels=reversed_attention_head_dim[i],
-                dual_cross_attention=dual_cross_attention,
-                use_linear_projection=use_linear_projection,
-                only_cross_attention=only_cross_attention[i],
-                upcast_attention=upcast_attention,
-                resnet_time_scale_shift=resnet_time_scale_shift,
-                use_inflated_groupnorm=use_inflated_groupnorm,
-                use_motion_module=use_motion_module and (res in motion_module_resolutions),
-                motion_module_type=motion_module_type,
-                motion_module_kwargs=motion_module_kwargs,
-                add_audio_layer=add_audio_layer,
-            )
-            self.up_blocks.append(up_block)
-            prev_output_channel = output_channel
-
-        # out
-        if use_inflated_groupnorm:
-            self.conv_norm_out = InflatedGroupNorm(
-                num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=norm_eps
-            )
-        else:
-            self.conv_norm_out = nn.GroupNorm(
-                num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=norm_eps
-            )
-        self.conv_act = nn.SiLU()
-
-        self.conv_out = zero_module(InflatedConv3d(block_out_channels[0], out_channels, kernel_size=3, padding=1))
-
-    def set_attention_slice(self, slice_size):
-        r"""
-        Enable sliced attention computation.
-
-        When this option is enabled, the attention module will split the input tensor in slices, to compute attention
-        in several steps. This is useful to save some memory in exchange for a small speed decrease.
-
-        Args:
-            slice_size (`str` or `int` or `list(int)`, *optional*, defaults to `"auto"`):
-                When `"auto"`, halves the input to the attention heads, so attention will be computed in two steps. If
-                `"max"`, maxium amount of memory will be saved by running only one slice at a time. If a number is
-                provided, uses as many slices as `attention_head_dim // slice_size`. In this case, `attention_head_dim`
-                must be a multiple of `slice_size`.
-        """
-        sliceable_head_dims = []
-
-        def fn_recursive_retrieve_slicable_dims(module: torch.nn.Module):
-            if hasattr(module, "set_attention_slice"):
-                sliceable_head_dims.append(module.sliceable_head_dim)
-
-            for child in module.children():
-                fn_recursive_retrieve_slicable_dims(child)
-
-        # retrieve number of attention layers
-        for module in self.children():
-            fn_recursive_retrieve_slicable_dims(module)
-
-        num_slicable_layers = len(sliceable_head_dims)
-
-        if slice_size == "auto":
-            # half the attention head size is usually a good trade-off between
-            # speed and memory
-            slice_size = [dim // 2 for dim in sliceable_head_dims]
-        elif slice_size == "max":
-            # make smallest slice possible
-            slice_size = num_slicable_layers * [1]
-
-        slice_size = num_slicable_layers * [slice_size] if not isinstance(slice_size, list) else slice_size
-
-        if len(slice_size) != len(sliceable_head_dims):
-            raise ValueError(
-                f"You have provided {len(slice_size)}, but {self.config} has {len(sliceable_head_dims)} different"
-                f" attention layers. Make sure to match `len(slice_size)` to be {len(sliceable_head_dims)}."
-            )
-
-        for i in range(len(slice_size)):
-            size = slice_size[i]
-            dim = sliceable_head_dims[i]
-            if size is not None and size > dim:
-                raise ValueError(f"size {size} has to be smaller or equal to {dim}.")
-
-        # Recursively walk through all the children.
-        # Any children which exposes the set_attention_slice method
-        # gets the message
-        def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: List[int]):
-            if hasattr(module, "set_attention_slice"):
-                module.set_attention_slice(slice_size.pop())
-
-            for child in module.children():
-                fn_recursive_set_attention_slice(child, slice_size)
-
-        reversed_slice_size = list(reversed(slice_size))
-        for module in self.children():
-            fn_recursive_set_attention_slice(module, reversed_slice_size)
-
-    def _set_gradient_checkpointing(self, module, value=False):
-        if isinstance(module, (CrossAttnDownBlock3D, DownBlock3D, CrossAttnUpBlock3D, UpBlock3D)):
-            module.gradient_checkpointing = value
-
-    def forward(
-        self,
-        sample: torch.FloatTensor,
-        timestep: Union[torch.Tensor, float, int],
-        encoder_hidden_states: torch.Tensor = None,
-        class_labels: Optional[torch.Tensor] = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        # support controlnet
-        down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None,
-        mid_block_additional_residual: Optional[torch.Tensor] = None,
-        return_dict: bool = True,
-    ) -> Union[UNet3DConditionOutput, Tuple]:
-        r"""
-        Args:
-            sample (`torch.FloatTensor`): (batch, channel, height, width) noisy inputs tensor
-            timestep (`torch.FloatTensor` or `float` or `int`): (batch) timesteps
-            encoder_hidden_states (`torch.FloatTensor`): (batch, sequence_length, feature_dim) encoder hidden states
-            return_dict (`bool`, *optional*, defaults to `True`):
-                Whether or not to return a [`models.unet_2d_condition.UNet2DConditionOutput`] instead of a plain tuple.
-
-        Returns:
-            [`~models.unet_2d_condition.UNet2DConditionOutput`] or `tuple`:
-            [`~models.unet_2d_condition.UNet2DConditionOutput`] if `return_dict` is True, otherwise a `tuple`. When
-            returning a tuple, the first element is the sample tensor.
-        """
-        # By default samples have to be AT least a multiple of the overall upsampling factor.
-        # The overall upsampling factor is equal to 2 ** (# num of upsampling layears).
-        # However, the upsampling interpolation output size can be forced to fit any upsampling size
-        # on the fly if necessary.
-        default_overall_up_factor = 2**self.num_upsamplers
-
-        # upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor`
-        forward_upsample_size = False
-        upsample_size = None
-
-        if any(s % default_overall_up_factor != 0 for s in sample.shape[-2:]):
-            logger.info("Forward upsample size to force interpolation output size.")
-            forward_upsample_size = True
-
-        # prepare attention_mask
-        if attention_mask is not None:
-            attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0
-            attention_mask = attention_mask.unsqueeze(1)
-
-        # center input if necessary
-        if self.config.center_input_sample:
-            sample = 2 * sample - 1.0
-
-        # time
-        timesteps = timestep
-        if not torch.is_tensor(timesteps):
-            # This would be a good case for the `match` statement (Python 3.10+)
-            is_mps = sample.device.type == "mps"
-            if isinstance(timestep, float):
-                dtype = torch.float32 if is_mps else torch.float64
-            else:
-                dtype = torch.int32 if is_mps else torch.int64
-            timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device)
-        elif len(timesteps.shape) == 0:
-            timesteps = timesteps[None].to(sample.device)
-
-        # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
-        timesteps = timesteps.expand(sample.shape[0])
-
-        t_emb = self.time_proj(timesteps)
-
-        # timesteps does not contain any weights and will always return f32 tensors
-        # but time_embedding might actually be running in fp16. so we need to cast here.
-        # there might be better ways to encapsulate this.
-        t_emb = t_emb.to(dtype=self.dtype)
-        emb = self.time_embedding(t_emb)
-
-        if self.class_embedding is not None:
-            if class_labels is None:
-                raise ValueError("class_labels should be provided when num_class_embeds > 0")
-
-            if self.config.class_embed_type == "timestep":
-                class_labels = self.time_proj(class_labels)
-
-            class_emb = self.class_embedding(class_labels).to(dtype=self.dtype)
-            emb = emb + class_emb
-
-        # pre-process
-        sample = self.conv_in(sample)
-
-        # down
-        down_block_res_samples = (sample,)
-        for downsample_block in self.down_blocks:
-            if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention:
-                sample, res_samples = downsample_block(
-                    hidden_states=sample,
-                    temb=emb,
-                    encoder_hidden_states=encoder_hidden_states,
-                    attention_mask=attention_mask,
-                )
-            else:
-                sample, res_samples = downsample_block(
-                    hidden_states=sample, temb=emb, encoder_hidden_states=encoder_hidden_states
-                )
-
-            down_block_res_samples += res_samples
-
-        # support controlnet
-        down_block_res_samples = list(down_block_res_samples)
-        if down_block_additional_residuals is not None:
-            for i, down_block_additional_residual in enumerate(down_block_additional_residuals):
-                if down_block_additional_residual.dim() == 4:  # boardcast
-                    down_block_additional_residual = down_block_additional_residual.unsqueeze(2)
-                down_block_res_samples[i] = down_block_res_samples[i] + down_block_additional_residual
-
-        # mid
-        sample = self.mid_block(
-            sample, emb, encoder_hidden_states=encoder_hidden_states, attention_mask=attention_mask
-        )
-
-        # support controlnet
-        if mid_block_additional_residual is not None:
-            if mid_block_additional_residual.dim() == 4:  # boardcast
-                mid_block_additional_residual = mid_block_additional_residual.unsqueeze(2)
-            sample = sample + mid_block_additional_residual
-
-        # up
-        for i, upsample_block in enumerate(self.up_blocks):
-            is_final_block = i == len(self.up_blocks) - 1
-
-            res_samples = down_block_res_samples[-len(upsample_block.resnets) :]
-            down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]
-
-            # if we have not reached the final block and need to forward the
-            # upsample size, we do it here
-            if not is_final_block and forward_upsample_size:
-                upsample_size = down_block_res_samples[-1].shape[2:]
-
-            if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention:
-                sample = upsample_block(
-                    hidden_states=sample,
-                    temb=emb,
-                    res_hidden_states_tuple=res_samples,
-                    encoder_hidden_states=encoder_hidden_states,
-                    upsample_size=upsample_size,
-                    attention_mask=attention_mask,
-                )
-            else:
-                sample = upsample_block(
-                    hidden_states=sample,
-                    temb=emb,
-                    res_hidden_states_tuple=res_samples,
-                    upsample_size=upsample_size,
-                    encoder_hidden_states=encoder_hidden_states,
-                )
-
-        # post-process
-        sample = self.conv_norm_out(sample)
-        sample = self.conv_act(sample)
-        sample = self.conv_out(sample)
-
-        if not return_dict:
-            return (sample,)
-
-        return UNet3DConditionOutput(sample=sample)
-
-    def load_state_dict(self, state_dict, strict=True):
-        # If the loaded checkpoint's in_channels or out_channels are different from config
-        if state_dict["conv_in.weight"].shape[1] != self.config.in_channels:
-            del state_dict["conv_in.weight"]
-            del state_dict["conv_in.bias"]
-        if state_dict["conv_out.weight"].shape[0] != self.config.out_channels:
-            del state_dict["conv_out.weight"]
-            del state_dict["conv_out.bias"]
-
-        # If the loaded checkpoint's cross_attention_dim is different from config
-        keys_to_remove = []
-        for key in state_dict:
-            if "attn2.to_k." in key or "attn2.to_v." in key:
-                if state_dict[key].shape[1] != self.config.cross_attention_dim:
-                    keys_to_remove.append(key)
-
-        for key in keys_to_remove:
-            del state_dict[key]
-
-        return super().load_state_dict(state_dict=state_dict, strict=strict)
-
-    @classmethod
-    def from_pretrained(cls, model_config: dict, ckpt_path: str, device="cpu"):
-        unet = cls.from_config(model_config).to(device)
-        if ckpt_path != "":
-            zero_rank_log(logger, f"Load from checkpoint: {ckpt_path}")
-            ckpt = torch.load(ckpt_path, map_location=device, weights_only=True)
-            if "global_step" in ckpt:
-                zero_rank_log(logger, f"resume from global_step: {ckpt['global_step']}")
-                resume_global_step = ckpt["global_step"]
-            else:
-                resume_global_step = 0
-            unet.load_state_dict(ckpt["state_dict"], strict=False)
-
-            del ckpt
-            torch.cuda.empty_cache()
-        else:
-            resume_global_step = 0
-
-        return unet, resume_global_step

latentsync/models/unet_blocks.py

@@ -1,777 +0,0 @@
-# Adapted from https://github.com/guoyww/AnimateDiff/blob/main/animatediff/models/unet_blocks.py
-
-import torch
-from torch import nn
-
-from .attention import Transformer3DModel
-from .resnet import Downsample3D, ResnetBlock3D, Upsample3D
-from .motion_module import get_motion_module
-
-
-def get_down_block(
-    down_block_type,
-    num_layers,
-    in_channels,
-    out_channels,
-    temb_channels,
-    add_downsample,
-    resnet_eps,
-    resnet_act_fn,
-    attn_num_head_channels,
-    resnet_groups=None,
-    cross_attention_dim=None,
-    downsample_padding=None,
-    dual_cross_attention=False,
-    use_linear_projection=False,
-    only_cross_attention=False,
-    upcast_attention=False,
-    resnet_time_scale_shift="default",
-    use_inflated_groupnorm=False,
-    use_motion_module=None,
-    motion_module_type=None,
-    motion_module_kwargs=None,
-    add_audio_layer=False,
-):
-    down_block_type = down_block_type[7:] if down_block_type.startswith("UNetRes") else down_block_type
-    if down_block_type == "DownBlock3D":
-        return DownBlock3D(
-            num_layers=num_layers,
-            in_channels=in_channels,
-            out_channels=out_channels,
-            temb_channels=temb_channels,
-            add_downsample=add_downsample,
-            resnet_eps=resnet_eps,
-            resnet_act_fn=resnet_act_fn,
-            resnet_groups=resnet_groups,
-            downsample_padding=downsample_padding,
-            resnet_time_scale_shift=resnet_time_scale_shift,
-            use_inflated_groupnorm=use_inflated_groupnorm,
-            use_motion_module=use_motion_module,
-            motion_module_type=motion_module_type,
-            motion_module_kwargs=motion_module_kwargs,
-        )
-    elif down_block_type == "CrossAttnDownBlock3D":
-        if cross_attention_dim is None:
-            raise ValueError("cross_attention_dim must be specified for CrossAttnDownBlock3D")
-        return CrossAttnDownBlock3D(
-            num_layers=num_layers,
-            in_channels=in_channels,
-            out_channels=out_channels,
-            temb_channels=temb_channels,
-            add_downsample=add_downsample,
-            resnet_eps=resnet_eps,
-            resnet_act_fn=resnet_act_fn,
-            resnet_groups=resnet_groups,
-            downsample_padding=downsample_padding,
-            cross_attention_dim=cross_attention_dim,
-            attn_num_head_channels=attn_num_head_channels,
-            dual_cross_attention=dual_cross_attention,
-            use_linear_projection=use_linear_projection,
-            only_cross_attention=only_cross_attention,
-            upcast_attention=upcast_attention,
-            resnet_time_scale_shift=resnet_time_scale_shift,
-            use_inflated_groupnorm=use_inflated_groupnorm,
-            use_motion_module=use_motion_module,
-            motion_module_type=motion_module_type,
-            motion_module_kwargs=motion_module_kwargs,
-            add_audio_layer=add_audio_layer,
-        )
-    raise ValueError(f"{down_block_type} does not exist.")
-
-
-def get_up_block(
-    up_block_type,
-    num_layers,
-    in_channels,
-    out_channels,
-    prev_output_channel,
-    temb_channels,
-    add_upsample,
-    resnet_eps,
-    resnet_act_fn,
-    attn_num_head_channels,
-    resnet_groups=None,
-    cross_attention_dim=None,
-    dual_cross_attention=False,
-    use_linear_projection=False,
-    only_cross_attention=False,
-    upcast_attention=False,
-    resnet_time_scale_shift="default",
-    use_inflated_groupnorm=False,
-    use_motion_module=None,
-    motion_module_type=None,
-    motion_module_kwargs=None,
-    add_audio_layer=False,
-):
-    up_block_type = up_block_type[7:] if up_block_type.startswith("UNetRes") else up_block_type
-    if up_block_type == "UpBlock3D":
-        return UpBlock3D(
-            num_layers=num_layers,
-            in_channels=in_channels,
-            out_channels=out_channels,
-            prev_output_channel=prev_output_channel,
-            temb_channels=temb_channels,
-            add_upsample=add_upsample,
-            resnet_eps=resnet_eps,
-            resnet_act_fn=resnet_act_fn,
-            resnet_groups=resnet_groups,
-            resnet_time_scale_shift=resnet_time_scale_shift,
-            use_inflated_groupnorm=use_inflated_groupnorm,
-            use_motion_module=use_motion_module,
-            motion_module_type=motion_module_type,
-            motion_module_kwargs=motion_module_kwargs,
-        )
-    elif up_block_type == "CrossAttnUpBlock3D":
-        if cross_attention_dim is None:
-            raise ValueError("cross_attention_dim must be specified for CrossAttnUpBlock3D")
-        return CrossAttnUpBlock3D(
-            num_layers=num_layers,
-            in_channels=in_channels,
-            out_channels=out_channels,
-            prev_output_channel=prev_output_channel,
-            temb_channels=temb_channels,
-            add_upsample=add_upsample,
-            resnet_eps=resnet_eps,
-            resnet_act_fn=resnet_act_fn,
-            resnet_groups=resnet_groups,
-            cross_attention_dim=cross_attention_dim,
-            attn_num_head_channels=attn_num_head_channels,
-            dual_cross_attention=dual_cross_attention,
-            use_linear_projection=use_linear_projection,
-            only_cross_attention=only_cross_attention,
-            upcast_attention=upcast_attention,
-            resnet_time_scale_shift=resnet_time_scale_shift,
-            use_inflated_groupnorm=use_inflated_groupnorm,
-            use_motion_module=use_motion_module,
-            motion_module_type=motion_module_type,
-            motion_module_kwargs=motion_module_kwargs,
-            add_audio_layer=add_audio_layer,
-        )
-    raise ValueError(f"{up_block_type} does not exist.")
-
-
-class UNetMidBlock3DCrossAttn(nn.Module):
-    def __init__(
-        self,
-        in_channels: int,
-        temb_channels: int,
-        dropout: float = 0.0,
-        num_layers: int = 1,
-        resnet_eps: float = 1e-6,
-        resnet_time_scale_shift: str = "default",
-        resnet_act_fn: str = "swish",
-        resnet_groups: int = 32,
-        resnet_pre_norm: bool = True,
-        attn_num_head_channels=1,
-        output_scale_factor=1.0,
-        cross_attention_dim=1280,
-        dual_cross_attention=False,
-        use_linear_projection=False,
-        upcast_attention=False,
-        use_inflated_groupnorm=False,
-        use_motion_module=None,
-        motion_module_type=None,
-        motion_module_kwargs=None,
-        add_audio_layer=False,
-    ):
-        super().__init__()
-
-        self.has_cross_attention = True
-        self.attn_num_head_channels = attn_num_head_channels
-        resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)
-
-        # there is always at least one resnet
-        resnets = [
-            ResnetBlock3D(
-                in_channels=in_channels,
-                out_channels=in_channels,
-                temb_channels=temb_channels,
-                eps=resnet_eps,
-                groups=resnet_groups,
-                dropout=dropout,
-                time_embedding_norm=resnet_time_scale_shift,
-                non_linearity=resnet_act_fn,
-                output_scale_factor=output_scale_factor,
-                pre_norm=resnet_pre_norm,
-                use_inflated_groupnorm=use_inflated_groupnorm,
-            )
-        ]
-        attentions = []
-        motion_modules = []
-
-        for _ in range(num_layers):
-            if dual_cross_attention:
-                raise NotImplementedError
-            attentions.append(
-                Transformer3DModel(
-                    attn_num_head_channels,
-                    in_channels // attn_num_head_channels,
-                    in_channels=in_channels,
-                    num_layers=1,
-                    cross_attention_dim=cross_attention_dim,
-                    norm_num_groups=resnet_groups,
-                    use_linear_projection=use_linear_projection,
-                    upcast_attention=upcast_attention,
-                    add_audio_layer=add_audio_layer,
-                )
-            )
-            motion_modules.append(
-                get_motion_module(
-                    in_channels=in_channels,
-                    motion_module_type=motion_module_type,
-                    motion_module_kwargs=motion_module_kwargs,
-                )
-                if use_motion_module
-                else None
-            )
-            resnets.append(
-                ResnetBlock3D(
-                    in_channels=in_channels,
-                    out_channels=in_channels,
-                    temb_channels=temb_channels,
-                    eps=resnet_eps,
-                    groups=resnet_groups,
-                    dropout=dropout,
-                    time_embedding_norm=resnet_time_scale_shift,
-                    non_linearity=resnet_act_fn,
-                    output_scale_factor=output_scale_factor,
-                    pre_norm=resnet_pre_norm,
-                    use_inflated_groupnorm=use_inflated_groupnorm,
-                )
-            )
-
-        self.attentions = nn.ModuleList(attentions)
-        self.resnets = nn.ModuleList(resnets)
-        self.motion_modules = nn.ModuleList(motion_modules)
-
-    def forward(self, hidden_states, temb=None, encoder_hidden_states=None, attention_mask=None):
-        hidden_states = self.resnets[0](hidden_states, temb)
-        for attn, resnet, motion_module in zip(self.attentions, self.resnets[1:], self.motion_modules):
-            hidden_states = attn(
-                hidden_states,
-                encoder_hidden_states=encoder_hidden_states,
-                return_dict=False,
-            )[0]
-
-            if motion_module is not None:
-                hidden_states = motion_module(hidden_states, temb, encoder_hidden_states=encoder_hidden_states)
-            hidden_states = resnet(hidden_states, temb)
-
-        return hidden_states
-
-
-class CrossAttnDownBlock3D(nn.Module):
-    def __init__(
-        self,
-        in_channels: int,
-        out_channels: int,
-        temb_channels: int,
-        dropout: float = 0.0,
-        num_layers: int = 1,
-        resnet_eps: float = 1e-6,
-        resnet_time_scale_shift: str = "default",
-        resnet_act_fn: str = "swish",
-        resnet_groups: int = 32,
-        resnet_pre_norm: bool = True,
-        attn_num_head_channels=1,
-        cross_attention_dim=1280,
-        output_scale_factor=1.0,
-        downsample_padding=1,
-        add_downsample=True,
-        dual_cross_attention=False,
-        use_linear_projection=False,
-        only_cross_attention=False,
-        upcast_attention=False,
-        use_inflated_groupnorm=False,
-        use_motion_module=None,
-        motion_module_type=None,
-        motion_module_kwargs=None,
-        add_audio_layer=False,
-    ):
-        super().__init__()
-        resnets = []
-        attentions = []
-        motion_modules = []
-
-        self.has_cross_attention = True
-        self.attn_num_head_channels = attn_num_head_channels
-
-        for i in range(num_layers):
-            in_channels = in_channels if i == 0 else out_channels
-            resnets.append(
-                ResnetBlock3D(
-                    in_channels=in_channels,
-                    out_channels=out_channels,
-                    temb_channels=temb_channels,
-                    eps=resnet_eps,
-                    groups=resnet_groups,
-                    dropout=dropout,
-                    time_embedding_norm=resnet_time_scale_shift,
-                    non_linearity=resnet_act_fn,
-                    output_scale_factor=output_scale_factor,
-                    pre_norm=resnet_pre_norm,
-                    use_inflated_groupnorm=use_inflated_groupnorm,
-                )
-            )
-            if dual_cross_attention:
-                raise NotImplementedError
-            attentions.append(
-                Transformer3DModel(
-                    attn_num_head_channels,
-                    out_channels // attn_num_head_channels,
-                    in_channels=out_channels,
-                    num_layers=1,
-                    cross_attention_dim=cross_attention_dim,
-                    norm_num_groups=resnet_groups,
-                    use_linear_projection=use_linear_projection,
-                    only_cross_attention=only_cross_attention,
-                    upcast_attention=upcast_attention,
-                    add_audio_layer=add_audio_layer,
-                )
-            )
-            motion_modules.append(
-                get_motion_module(
-                    in_channels=out_channels,
-                    motion_module_type=motion_module_type,
-                    motion_module_kwargs=motion_module_kwargs,
-                )
-                if use_motion_module
-                else None
-            )
-
-        self.attentions = nn.ModuleList(attentions)
-        self.resnets = nn.ModuleList(resnets)
-        self.motion_modules = nn.ModuleList(motion_modules)
-
-        if add_downsample:
-            self.downsamplers = nn.ModuleList(
-                [
-                    Downsample3D(
-                        out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op"
-                    )
-                ]
-            )
-        else:
-            self.downsamplers = None
-
-        self.gradient_checkpointing = False
-
-    def forward(self, hidden_states, temb=None, encoder_hidden_states=None, attention_mask=None):
-        output_states = ()
-
-        for resnet, attn, motion_module in zip(self.resnets, self.attentions, self.motion_modules):
-            if torch.is_grad_enabled() and self.gradient_checkpointing:
-
-                def create_custom_forward(module, return_dict=None):
-                    def custom_forward(*inputs):
-                        if return_dict is not None:
-                            return module(*inputs, return_dict=return_dict)
-                        else:
-                            return module(*inputs)
-
-                    return custom_forward
-
-                hidden_states = torch.utils.checkpoint.checkpoint(
-                    create_custom_forward(resnet), hidden_states, temb, use_reentrant=False
-                )
-                hidden_states = torch.utils.checkpoint.checkpoint(
-                    create_custom_forward(attn, return_dict=False),
-                    hidden_states,
-                    encoder_hidden_states,
-                    use_reentrant=False,
-                )[0]
-
-                if motion_module is not None:
-                    hidden_states = torch.utils.checkpoint.checkpoint(
-                        create_custom_forward(motion_module),
-                        hidden_states,
-                        temb,
-                        encoder_hidden_states,
-                        use_reentrant=False,
-                    )
-            else:
-                hidden_states = resnet(hidden_states, temb)
-                hidden_states = attn(hidden_states, encoder_hidden_states=encoder_hidden_states).sample
-
-                if motion_module is not None:
-                    hidden_states = motion_module(hidden_states, temb, encoder_hidden_states=encoder_hidden_states)
-
-            output_states += (hidden_states,)
-
-        if self.downsamplers is not None:
-            for downsampler in self.downsamplers:
-                hidden_states = downsampler(hidden_states)
-
-            output_states += (hidden_states,)
-
-        return hidden_states, output_states
-
-
-class DownBlock3D(nn.Module):
-    def __init__(
-        self,
-        in_channels: int,
-        out_channels: int,
-        temb_channels: int,
-        dropout: float = 0.0,
-        num_layers: int = 1,
-        resnet_eps: float = 1e-6,
-        resnet_time_scale_shift: str = "default",
-        resnet_act_fn: str = "swish",
-        resnet_groups: int = 32,
-        resnet_pre_norm: bool = True,
-        output_scale_factor=1.0,
-        add_downsample=True,
-        downsample_padding=1,
-        use_inflated_groupnorm=False,
-        use_motion_module=None,
-        motion_module_type=None,
-        motion_module_kwargs=None,
-    ):
-        super().__init__()
-        resnets = []
-        motion_modules = []
-
-        for i in range(num_layers):
-            in_channels = in_channels if i == 0 else out_channels
-            resnets.append(
-                ResnetBlock3D(
-                    in_channels=in_channels,
-                    out_channels=out_channels,
-                    temb_channels=temb_channels,
-                    eps=resnet_eps,
-                    groups=resnet_groups,
-                    dropout=dropout,
-                    time_embedding_norm=resnet_time_scale_shift,
-                    non_linearity=resnet_act_fn,
-                    output_scale_factor=output_scale_factor,
-                    pre_norm=resnet_pre_norm,
-                    use_inflated_groupnorm=use_inflated_groupnorm,
-                )
-            )
-            motion_modules.append(
-                get_motion_module(
-                    in_channels=out_channels,
-                    motion_module_type=motion_module_type,
-                    motion_module_kwargs=motion_module_kwargs,
-                )
-                if use_motion_module
-                else None
-            )
-
-        self.resnets = nn.ModuleList(resnets)
-        self.motion_modules = nn.ModuleList(motion_modules)
-
-        if add_downsample:
-            self.downsamplers = nn.ModuleList(
-                [
-                    Downsample3D(
-                        out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op"
-                    )
-                ]
-            )
-        else:
-            self.downsamplers = None
-
-        self.gradient_checkpointing = False
-
-    def forward(self, hidden_states, temb=None, encoder_hidden_states=None):
-        output_states = ()
-
-        for resnet, motion_module in zip(self.resnets, self.motion_modules):
-            if torch.is_grad_enabled() and self.gradient_checkpointing:
-
-                def create_custom_forward(module):
-                    def custom_forward(*inputs):
-                        return module(*inputs)
-
-                    return custom_forward
-
-                hidden_states = torch.utils.checkpoint.checkpoint(
-                    create_custom_forward(resnet), hidden_states, temb, use_reentrant=False
-                )
-
-                if motion_module is not None:
-                    hidden_states = torch.utils.checkpoint.checkpoint(
-                        create_custom_forward(motion_module),
-                        hidden_states,
-                        temb,
-                        encoder_hidden_states,
-                        use_reentrant=False,
-                    )
-            else:
-                hidden_states = resnet(hidden_states, temb)
-
-                if motion_module is not None:
-                    hidden_states = motion_module(hidden_states, temb, encoder_hidden_states=encoder_hidden_states)
-
-            output_states += (hidden_states,)
-
-        if self.downsamplers is not None:
-            for downsampler in self.downsamplers:
-                hidden_states = downsampler(hidden_states)
-
-            output_states += (hidden_states,)
-
-        return hidden_states, output_states
-
-
-class CrossAttnUpBlock3D(nn.Module):
-    def __init__(
-        self,
-        in_channels: int,
-        out_channels: int,
-        prev_output_channel: int,
-        temb_channels: int,
-        dropout: float = 0.0,
-        num_layers: int = 1,
-        resnet_eps: float = 1e-6,
-        resnet_time_scale_shift: str = "default",
-        resnet_act_fn: str = "swish",
-        resnet_groups: int = 32,
-        resnet_pre_norm: bool = True,
-        attn_num_head_channels=1,
-        cross_attention_dim=1280,
-        output_scale_factor=1.0,
-        add_upsample=True,
-        dual_cross_attention=False,
-        use_linear_projection=False,
-        only_cross_attention=False,
-        upcast_attention=False,
-        use_inflated_groupnorm=False,
-        use_motion_module=None,
-        motion_module_type=None,
-        motion_module_kwargs=None,
-        add_audio_layer=False,
-    ):
-        super().__init__()
-        resnets = []
-        attentions = []
-        motion_modules = []
-
-        self.has_cross_attention = True
-        self.attn_num_head_channels = attn_num_head_channels
-
-        for i in range(num_layers):
-            res_skip_channels = in_channels if (i == num_layers - 1) else out_channels
-            resnet_in_channels = prev_output_channel if i == 0 else out_channels
-
-            resnets.append(
-                ResnetBlock3D(
-                    in_channels=resnet_in_channels + res_skip_channels,
-                    out_channels=out_channels,
-                    temb_channels=temb_channels,
-                    eps=resnet_eps,
-                    groups=resnet_groups,
-                    dropout=dropout,
-                    time_embedding_norm=resnet_time_scale_shift,
-                    non_linearity=resnet_act_fn,
-                    output_scale_factor=output_scale_factor,
-                    pre_norm=resnet_pre_norm,
-                    use_inflated_groupnorm=use_inflated_groupnorm,
-                )
-            )
-            if dual_cross_attention:
-                raise NotImplementedError
-            attentions.append(
-                Transformer3DModel(
-                    attn_num_head_channels,
-                    out_channels // attn_num_head_channels,
-                    in_channels=out_channels,
-                    num_layers=1,
-                    cross_attention_dim=cross_attention_dim,
-                    norm_num_groups=resnet_groups,
-                    use_linear_projection=use_linear_projection,
-                    only_cross_attention=only_cross_attention,
-                    upcast_attention=upcast_attention,
-                    add_audio_layer=add_audio_layer,
-                )
-            )
-            motion_modules.append(
-                get_motion_module(
-                    in_channels=out_channels,
-                    motion_module_type=motion_module_type,
-                    motion_module_kwargs=motion_module_kwargs,
-                )
-                if use_motion_module
-                else None
-            )
-
-        self.attentions = nn.ModuleList(attentions)
-        self.resnets = nn.ModuleList(resnets)
-        self.motion_modules = nn.ModuleList(motion_modules)
-
-        if add_upsample:
-            self.upsamplers = nn.ModuleList([Upsample3D(out_channels, use_conv=True, out_channels=out_channels)])
-        else:
-            self.upsamplers = None
-
-        self.gradient_checkpointing = False
-
-    def forward(
-        self,
-        hidden_states,
-        res_hidden_states_tuple,
-        temb=None,
-        encoder_hidden_states=None,
-        upsample_size=None,
-        attention_mask=None,
-    ):
-        for resnet, attn, motion_module in zip(self.resnets, self.attentions, self.motion_modules):
-            # pop res hidden states
-            res_hidden_states = res_hidden_states_tuple[-1]
-            res_hidden_states_tuple = res_hidden_states_tuple[:-1]
-            hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
-
-            if torch.is_grad_enabled() and self.gradient_checkpointing:
-
-                def create_custom_forward(module, return_dict=None):
-                    def custom_forward(*inputs):
-                        if return_dict is not None:
-                            return module(*inputs, return_dict=return_dict)
-                        else:
-                            return module(*inputs)
-
-                    return custom_forward
-
-                hidden_states = torch.utils.checkpoint.checkpoint(
-                    create_custom_forward(resnet), hidden_states, temb, use_reentrant=False
-                )
-                hidden_states = torch.utils.checkpoint.checkpoint(
-                    create_custom_forward(attn, return_dict=False),
-                    hidden_states,
-                    encoder_hidden_states,
-                    use_reentrant=False,
-                )[0]
-
-                if motion_module is not None:
-                    hidden_states = torch.utils.checkpoint.checkpoint(
-                        create_custom_forward(motion_module),
-                        hidden_states,
-                        temb,
-                        encoder_hidden_states,
-                        use_reentrant=False,
-                    )
-            else:
-                hidden_states = resnet(hidden_states, temb)
-                hidden_states = attn(hidden_states, encoder_hidden_states=encoder_hidden_states).sample
-
-                if motion_module is not None:
-                    hidden_states = motion_module(hidden_states, temb, encoder_hidden_states=encoder_hidden_states)
-
-        if self.upsamplers is not None:
-            for upsampler in self.upsamplers:
-                hidden_states = upsampler(hidden_states, upsample_size)
-
-        return hidden_states
-
-
-class UpBlock3D(nn.Module):
-    def __init__(
-        self,
-        in_channels: int,
-        prev_output_channel: int,
-        out_channels: int,
-        temb_channels: int,
-        dropout: float = 0.0,
-        num_layers: int = 1,
-        resnet_eps: float = 1e-6,
-        resnet_time_scale_shift: str = "default",
-        resnet_act_fn: str = "swish",
-        resnet_groups: int = 32,
-        resnet_pre_norm: bool = True,
-        output_scale_factor=1.0,
-        add_upsample=True,
-        use_inflated_groupnorm=False,
-        use_motion_module=None,
-        motion_module_type=None,
-        motion_module_kwargs=None,
-    ):
-        super().__init__()
-        resnets = []
-        motion_modules = []
-
-        for i in range(num_layers):
-            res_skip_channels = in_channels if (i == num_layers - 1) else out_channels
-            resnet_in_channels = prev_output_channel if i == 0 else out_channels
-
-            resnets.append(
-                ResnetBlock3D(
-                    in_channels=resnet_in_channels + res_skip_channels,
-                    out_channels=out_channels,
-                    temb_channels=temb_channels,
-                    eps=resnet_eps,
-                    groups=resnet_groups,
-                    dropout=dropout,
-                    time_embedding_norm=resnet_time_scale_shift,
-                    non_linearity=resnet_act_fn,
-                    output_scale_factor=output_scale_factor,
-                    pre_norm=resnet_pre_norm,
-                    use_inflated_groupnorm=use_inflated_groupnorm,
-                )
-            )
-            motion_modules.append(
-                get_motion_module(
-                    in_channels=out_channels,
-                    motion_module_type=motion_module_type,
-                    motion_module_kwargs=motion_module_kwargs,
-                )
-                if use_motion_module
-                else None
-            )
-
-        self.resnets = nn.ModuleList(resnets)
-        self.motion_modules = nn.ModuleList(motion_modules)
-
-        if add_upsample:
-            self.upsamplers = nn.ModuleList([Upsample3D(out_channels, use_conv=True, out_channels=out_channels)])
-        else:
-            self.upsamplers = None
-
-        self.gradient_checkpointing = False
-
-    def forward(
-        self,
-        hidden_states,
-        res_hidden_states_tuple,
-        temb=None,
-        upsample_size=None,
-        encoder_hidden_states=None,
-    ):
-        for resnet, motion_module in zip(self.resnets, self.motion_modules):
-            # pop res hidden states
-            res_hidden_states = res_hidden_states_tuple[-1]
-            res_hidden_states_tuple = res_hidden_states_tuple[:-1]
-            hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
-
-            if torch.is_grad_enabled() and self.gradient_checkpointing:
-
-                def create_custom_forward(module):
-                    def custom_forward(*inputs):
-                        return module(*inputs)
-
-                    return custom_forward
-
-                hidden_states = torch.utils.checkpoint.checkpoint(
-                    create_custom_forward(resnet), hidden_states, temb, use_reentrant=False
-                )
-
-                if motion_module is not None:
-                    hidden_states = torch.utils.checkpoint.checkpoint(
-                        create_custom_forward(motion_module),
-                        hidden_states,
-                        temb,
-                        encoder_hidden_states,
-                        use_reentrant=False,
-                    )
-            else:
-                hidden_states = resnet(hidden_states, temb)
-
-                if motion_module is not None:
-                    hidden_states = motion_module(hidden_states, temb, encoder_hidden_states=encoder_hidden_states)
-
-        if self.upsamplers is not None:
-            for upsampler in self.upsamplers:
-                hidden_states = upsampler(hidden_states, upsample_size)
-
-        return hidden_states

latentsync/models/utils.py

@@ -1,19 +0,0 @@
-# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-def zero_module(module):
-    # Zero out the parameters of a module and return it.
-    for p in module.parameters():
-        p.detach().zero_()
-    return module