Difix3D

You can use the DifixPipeline via the diffusers library for single-step diffusion-based artifact removal.

Standard Usage: Load the pretrained model nvidia/difix and pass an input image with a prompt (e.g., "remove degradation").

With Reference Image: To guide the denoising process, use the nvidia/difix_ref model and provide a ref_image to the pipeline call.

from pipeline_difix import DifixPipeline
from diffusers.utils import load_image

# Standard usage
pipe = DifixPipeline.from_pretrained("nvidia/difix", trust_remote_code=True)
pipe.to("cuda")

input_image = load_image("assets/example_input.png")
prompt = "remove degradation"

output_image = pipe(prompt, image=input_image, num_inference_steps=1, timesteps=[199], guidance_scale=0.0).images[0]
output_image.save("example_output.png")

Difix3D+ uses the Difix model as a post-processing step at render time to enhance novel views and recover sharp details that reconstruction methods might miss. Run inference on rendered images using src/inference_difix.py.

python src/inference_difix.py \
    --model_path "checkpoints/model.pkl" \
    --input_image "PATH_TO_IMAGES" \
    --prompt "remove degradation" \
    --output_dir "outputs/difix3d+" \
    --timestep 199

To set up the Difix3D environment, clone the repository and install the required dependencies using pip.

git clone https://github.com/nv-tlabs/Difix3D.git
cd Difix3D
pip install -r requirements.txt

To use Difix3D with gsplat, ensure gsplat is installed following its official instructions. Use the examples/gsplat/simple_trainer_difix3d.py script to run the finetuning process.

SCENE_ID=032dee9fb0a8bc1b90871dc5fe950080d0bcd3caf166447f44e60ca50ac04ec7
DATA=DATA_DIR/${SCENE_ID}/gaussian_splat
DATA_FACTOR=4
CKPT_PATH=CKPT_DIR/${SCENE_ID}/ckpts/ckpt_29999_rank0.pt
OUTPUT_DIR=outputs/difix3d/gsplat/${SCENE_ID}

CUDA_VISIBLE_DEVICES=0 python examples/gsplat/simple_trainer_difix3d.py default \
    --data_dir ${DATA} --data_factor ${DATA_FACTOR} \
    --result_dir ${OUTPUT_DIR} --no-normalize-world-space --test_every 1 --ckpt ${CKPT_PATH}

When using certain noise schedules (like zero terminal SNR), the guidance_scale can cause overexposure in generated images. You can mitigate this by providing a guidance_rescale value to the __call__ method. This factor mixes the guidance results with the original results to maintain better image quality.

Parameter:

• guidance_rescale (float): A factor used to rescale the noise prediction based on the guidance scale. Defaults to 0.0 (no rescaling).

To run inference on a single image using a trained checkpoint, place your model_*.pkl in the checkpoints directory and execute src/inference_difix.py.

python src/inference_difix.py \
    --model_path "checkpoints/model.pkl" \
    --input_image "assets/example_input.png" \
    --prompt "remove degradation" \
    --output_dir "outputs/difix" \
    --timestep 199

Datasets for training the single-step Difix model must be provided in a JSON format containing train and test splits. Each entry in a split is keyed by a {data_id} and contains paths to the input image, target image, reference image, and the associated prompt.

{
  "train": {
    "{data_id}": {
      "image": "{PATH_TO_IMAGE}",
      "target_image": "{PATH_TO_TARGET_IMAGE}",
      "ref_image": "{PATH_TO_REF_IMAGE}",
      "prompt": "remove degradation"
    }
  },
  "test": {
    "{data_id}": {
      "image": "{PATH_TO_IMAGE}",
      "target_image": "{PATH_TO_TARGET_IMAGE}",
      "ref_image": "{PATH_TO_REF_IMAGE}",
      "prompt": "remove degradation"
    }
  }
}

Difix3D requires a specific directory structure for scenes. Each {SCENE_ID} must contain a colmap directory (with sparse reconstruction files) and an images directory, along with multi-scale image directories (images_2, images_4, images_8).

DATA_DIR/
├── {SCENE_ID}
│   ├── colmap
│   │   ├── sparse
│   │   │   └── 0
│   │   │       ├── cameras.bin
│   │   │       ├── database.db
│   │   │       └── ...
│   ├── images
│   │   ├── image_train_000001.png
│   │   ├── image_train_000002.png
│   │   ├── ...
│   │   ├── image_eval_000200.png
│   │   ├── image_eval_000201.png
│   │   └── ...
│   ├── images_2
│   ├── images_4
│   └── images_8

The Difix class is the primary entry point for the model. It initializes a diffusion-based model using sd-turbo components (tokenizer, text encoder, VAE, and UNet). It supports loading from a pretrained path or initializing with random weights and adding LoRA adapters to the VAE for skip-connection learning.

Key arguments:

• pretrained_name: Name of the pretrained model (e.g., from Hugging Face).
• pretrained_path: Local path to a checkpoint file. If provided, it loads the UNet, VAE (with LoRA), and optimizer states.
• ckpt_folder: Directory for checkpoints.
• lora_rank_vae: The rank r for the VAE LoRA adapter (default is 4).
• mv_unet: If True, uses mv_unet.UNet2DConditionModel instead of the standard diffusers UNet.
• timestep: The diffusion timestep to use (default is 999).