README.md

---
license: apache-2.0
base_model: Wan-AI/Wan2.2-TI2V-5B-Diffusers
library_name: diffusers
tags:
  - vae
  - video
  - wan
  - world-model
---

# wan-vae-minecraft-352-640

Finetuned [`AutoencoderKLWan`](https://huggingface.co/docs/diffusers/main/en/api/models/autoencoderkl_wan)
derived from [`Wan-AI/Wan2.2-TI2V-5B-Diffusers`](https://huggingface.co/Wan-AI/Wan2.2-TI2V-5B-Diffusers).

## Usage

```python
from diffusers import AutoencoderKLWan
import torch

vae = AutoencoderKLWan.from_pretrained("aidanscasnnell/wan-vae-minecraft-352-640", torch_dtype=torch.float32)
```

## Dataset

- **Dataset:** Minecraft (OpenAI VPT)
- **Training resolution:** 352x640 (H×W)

Frames sampled from the Minecraft gameplay videos released by OpenAI's Video Pre-Training project ([Video-Pre-Training](https://github.com/openai/Video-Pre-Training)). Episodes are RGB video; the VAE is trained to reconstruct them independently of actions.

## Training

- Last saved step: `51000`

### Reproduction
- Training script: `scripts/finetune_wan_vae.py`
- Source commit: `056f3a6172e738c84ef0c0922ac3531a3cf4fbe9`

### Best metrics
- `metric`: val/loss
- `mode`: min
- `value`: 1.06140953818957
- `step`: 51000

<details>
<summary>Training config</summary>

```json
{
  "data": {
    "dataset": "minerl",
    "seed": 123,
    "ctx_len_fr": 17,
    "pred_len_fr": 16,
    "eval_pred_len_fr": 60,
    "eval_data_stride": 1,
    "resize_resolution": [
      352,
      640
    ],
    "path": "1x-technologies/worldmodel_raw_data",
    "use_latents": false,
    "bfloat16_latents": true,
    "use_precomputed_index": true,
    "minerl_dir": "/mnt/minecraft",
    "tasks": [
      "all"
    ],
    "minerl_split_seed": 42,
    "minerl_drop_last": true,
    "minerl_pad_to_len": false,
    "minerl_latents_dir": null,
    "minerl_latents_subdir": "latents",
    "minerl_window_stride_lat": 1,
    "minerl_return_z0_latent": false,
    "minerl_total_val_clips": 512,
    "minerl_total_test_clips": 1024,
    "eval_window_mode": "prefix",
    "num_workers": 8,
    "persistent_workers": true,
    "prefetch_factor": 4,
    "pin_memory": true
  },
  "model": {
    "hf_id": "Wan-AI/Wan2.2-TI2V-5B-Diffusers",
    "hf_subfolder": "vae",
    "torch_dtype": "float32"
  },
  "loss": {
    "l1_w": 3.0,
    "kl_w": 3e-06,
    "lpips_w": 3.0,
    "temporal_w": 0.5,
    "use_lpips": true,
    "lpips_net": "vgg",
    "lpips_on_frames": true,
    "use_gan": false,
    "gan_w": 0.1,
    "disc_lr": 0.0002,
    "disc_steps_per_gen_step": 1,
    "disc_start_step": 0,
    "r1_gamma": 0.0
  },
  "train": {
    "finetune": "decoder",
    "seed": 123,
    "device": "cuda",
    "amp": true,
    "amp_dtype": "bfloat16",
    "batch_size": 1,
    "val_batch_size": 1,
    "grad_accum_steps": 2,
    "lr": 1e-05,
    "betas": [
      0.9,
      0.999
    ],
    "weight_decay": 0.0,
    "max_steps": 150000,
    "log_every": 50,
    "eval_every": 500,
    "save_every": 500,
    "video_every": 500,
    "video_fps": 8,
    "max_val_batches": 50,
    "best_metric_name": "val/loss",
    "best_metric_mode": "min",
    "resume_path": "/root/radit/outputs/2026-04-17/18-06-56/checkpoints/best.pt"
  },
  "compile": {
    "enabled": true,
    "backend": "inductor",
    "mode": "max-autotune",
    "fullgraph": false
  },
  "metrics": {
    "enabled": true,
    "metric_names": [
      "psnr",
      "ssim",
      "lpips"
    ],
    "metrics_num_samples": 256,
    "psnr_log_stride": 1,
    "max_val": 1.0
  },
  "logger": {
    "use_wandb": true,
    "run_name": "resumed-res_[352, 640]-all_data-temporal_w_0.5-updated-tasks_['all']-finetune_decoder-bsize_1-accum_2-lr_1e-05-pred_len_16",
    "project": "wan-vae-finetune"
  }
}
```
</details>


## License

Inherits the license of the base model (Wan-AI/Wan2.2-TI2V-5B-Diffusers); verify terms before redistribution.