Add files using upload-large-folder tool

demo.py

@@ -0,0 +1,631 @@
+"""
+Demo for Self-Forcing.
+"""
+
+import os
+import re
+import random
+import time
+import base64
+import argparse
+import hashlib
+import subprocess
+import urllib.request
+from io import BytesIO
+from PIL import Image
+import numpy as np
+import torch
+from omegaconf import OmegaConf
+from flask import Flask, render_template, jsonify
+from flask_socketio import SocketIO, emit
+import queue
+from threading import Thread, Event
+
+from pipeline import CausalInferencePipeline
+from demo_utils.constant import ZERO_VAE_CACHE
+from demo_utils.vae_block3 import VAEDecoderWrapper
+from utils.wan_wrapper import WanDiffusionWrapper, WanTextEncoder
+from demo_utils.utils import generate_timestamp
+from demo_utils.memory import gpu, get_cuda_free_memory_gb, DynamicSwapInstaller, move_model_to_device_with_memory_preservation
+
+# Parse arguments
+parser = argparse.ArgumentParser()
+parser.add_argument('--port', type=int, default=5001)
+parser.add_argument('--host', type=str, default='0.0.0.0')
+parser.add_argument("--checkpoint_path", type=str, default='./checkpoints/self_forcing_dmd.pt')
+parser.add_argument("--config_path", type=str, default='./configs/self_forcing_dmd.yaml')
+parser.add_argument('--trt', action='store_true')
+args = parser.parse_args()
+
+print(f'Free VRAM {get_cuda_free_memory_gb(gpu)} GB')
+low_memory = get_cuda_free_memory_gb(gpu) < 40
+
+# Load models
+config = OmegaConf.load(args.config_path)
+default_config = OmegaConf.load("configs/default_config.yaml")
+config = OmegaConf.merge(default_config, config)
+
+text_encoder = WanTextEncoder()
+
+# Global variables for dynamic model switching
+current_vae_decoder = None
+current_use_taehv = False
+fp8_applied = False
+torch_compile_applied = False
+global frame_number
+frame_number = 0
+anim_name = ""
+frame_rate = 6
+
+def initialize_vae_decoder(use_taehv=False, use_trt=False):
+    """Initialize VAE decoder based on the selected option"""
+    global current_vae_decoder, current_use_taehv
+
+    if use_trt:
+        from demo_utils.vae import VAETRTWrapper
+        current_vae_decoder = VAETRTWrapper()
+        return current_vae_decoder
+
+    if use_taehv:
+        from demo_utils.taehv import TAEHV
+        # Check if taew2_1.pth exists in checkpoints folder, download if missing
+        taehv_checkpoint_path = "checkpoints/taew2_1.pth"
+        if not os.path.exists(taehv_checkpoint_path):
+            print(f"taew2_1.pth not found in checkpoints folder {taehv_checkpoint_path}. Downloading...")
+            os.makedirs("checkpoints", exist_ok=True)
+            download_url = "https://github.com/madebyollin/taehv/raw/main/taew2_1.pth"
+            try:
+                urllib.request.urlretrieve(download_url, taehv_checkpoint_path)
+                print(f"Successfully downloaded taew2_1.pth to {taehv_checkpoint_path}")
+            except Exception as e:
+                print(f"Failed to download taew2_1.pth: {e}")
+                raise
+
+        class DotDict(dict):
+            __getattr__ = dict.__getitem__
+            __setattr__ = dict.__setitem__
+
+        class TAEHVDiffusersWrapper(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.dtype = torch.float16
+                self.taehv = TAEHV(checkpoint_path=taehv_checkpoint_path).to(self.dtype)
+                self.config = DotDict(scaling_factor=1.0)
+
+            def decode(self, latents, return_dict=None):
+                # n, c, t, h, w = latents.shape
+                # low-memory, set parallel=True for faster + higher memory
+                return self.taehv.decode_video(latents, parallel=False).mul_(2).sub_(1)
+
+        current_vae_decoder = TAEHVDiffusersWrapper()
+    else:
+        current_vae_decoder = VAEDecoderWrapper()
+        vae_state_dict = torch.load('wan_models/Wan2.1-T2V-1.3B/Wan2.1_VAE.pth', map_location="cpu")
+        decoder_state_dict = {}
+        for key, value in vae_state_dict.items():
+            if 'decoder.' in key or 'conv2' in key:
+                decoder_state_dict[key] = value
+        current_vae_decoder.load_state_dict(decoder_state_dict)
+
+    current_vae_decoder.eval()
+    current_vae_decoder.to(dtype=torch.float16)
+    current_vae_decoder.requires_grad_(False)
+    current_vae_decoder.to(gpu)
+    current_use_taehv = use_taehv
+
+    print(f"✅ VAE decoder initialized with {'TAEHV' if use_taehv else 'default VAE'}")
+    return current_vae_decoder
+
+
+# Initialize with default VAE
+vae_decoder = initialize_vae_decoder(use_taehv=False, use_trt=args.trt)
+
+transformer = WanDiffusionWrapper(is_causal=True)
+state_dict = torch.load(args.checkpoint_path, map_location="cpu")
+transformer.load_state_dict(state_dict['generator_ema'])
+
+text_encoder.eval()
+transformer.eval()
+
+transformer.to(dtype=torch.float16)
+text_encoder.to(dtype=torch.bfloat16)
+
+text_encoder.requires_grad_(False)
+transformer.requires_grad_(False)
+
+pipeline = CausalInferencePipeline(
+    config,
+    device=gpu,
+    generator=transformer,
+    text_encoder=text_encoder,
+    vae=vae_decoder
+)
+
+if low_memory:
+    DynamicSwapInstaller.install_model(text_encoder, device=gpu)
+else:
+    text_encoder.to(gpu)
+transformer.to(gpu)
+
+# Flask and SocketIO setup
+app = Flask(__name__)
+app.config['SECRET_KEY'] = 'frontend_buffered_demo'
+socketio = SocketIO(app, cors_allowed_origins="*")
+
+generation_active = False
+stop_event = Event()
+frame_send_queue = queue.Queue()
+sender_thread = None
+models_compiled = False
+
+
+def tensor_to_base64_frame(frame_tensor):
+    """Convert a single frame tensor to base64 image string."""
+    global frame_number, anim_name
+    # Clamp and normalize to 0-255
+    frame = torch.clamp(frame_tensor.float(), -1., 1.) * 127.5 + 127.5
+    frame = frame.to(torch.uint8).cpu().numpy()
+
+    # CHW -> HWC
+    if len(frame.shape) == 3:
+        frame = np.transpose(frame, (1, 2, 0))
+
+    # Convert to PIL Image
+    if frame.shape[2] == 3:  # RGB
+        image = Image.fromarray(frame, 'RGB')
+    else:  # Handle other formats
+        image = Image.fromarray(frame)
+
+    # Convert to base64
+    buffer = BytesIO()
+    image.save(buffer, format='JPEG', quality=100)
+    if not os.path.exists("./images/%s" % anim_name):
+        os.makedirs("./images/%s" % anim_name)
+    frame_number += 1
+    image.save("./images/%s/%s_%03d.jpg" % (anim_name, anim_name, frame_number))
+    img_str = base64.b64encode(buffer.getvalue()).decode()
+    return f"data:image/jpeg;base64,{img_str}"
+
+
+def frame_sender_worker():
+    """Background thread that processes frame send queue non-blocking."""
+    global frame_send_queue, generation_active, stop_event
+
+    print("📡 Frame sender thread started")
+
+    while True:
+        frame_data = None
+        try:
+            # Get frame data from queue
+            frame_data = frame_send_queue.get(timeout=1.0)
+
+            if frame_data is None:  # Shutdown signal
+                frame_send_queue.task_done()  # Mark shutdown signal as done
+                break
+
+            frame_tensor, frame_index, block_index, job_id = frame_data
+
+            # Convert tensor to base64
+            base64_frame = tensor_to_base64_frame(frame_tensor)
+
+            # Send via SocketIO
+            try:
+                socketio.emit('frame_ready', {
+                    'data': base64_frame,
+                    'frame_index': frame_index,
+                    'block_index': block_index,
+                    'job_id': job_id
+                })
+            except Exception as e:
+                print(f"⚠️ Failed to send frame {frame_index}: {e}")
+
+            frame_send_queue.task_done()
+
+        except queue.Empty:
+            # Check if we should continue running
+            if not generation_active and frame_send_queue.empty():
+                break
+        except Exception as e:
+            print(f"❌ Frame sender error: {e}")
+            # Make sure to mark task as done even if there's an error
+            if frame_data is not None:
+                try:
+                    frame_send_queue.task_done()
+                except Exception as e:
+                    print(f"❌ Failed to mark frame task as done: {e}")
+            break
+
+    print("📡 Frame sender thread stopped")
+
+
+@torch.no_grad()
+def generate_video_stream(prompt, seed, enable_torch_compile=False, enable_fp8=False, use_taehv=False):
+    """Generate video and push frames immediately to frontend."""
+    global generation_active, stop_event, frame_send_queue, sender_thread, models_compiled, torch_compile_applied, fp8_applied, current_vae_decoder, current_use_taehv, frame_rate, anim_name
+
+    try:
+        generation_active = True
+        stop_event.clear()
+        job_id = generate_timestamp()
+
+        # Start frame sender thread if not already running
+        if sender_thread is None or not sender_thread.is_alive():
+            sender_thread = Thread(target=frame_sender_worker, daemon=True)
+            sender_thread.start()
+
+        # Emit progress updates
+        def emit_progress(message, progress):
+            try:
+                socketio.emit('progress', {
+                    'message': message,
+                    'progress': progress,
+                    'job_id': job_id
+                })
+            except Exception as e:
+                print(f"❌ Failed to emit progress: {e}")
+
+        emit_progress('Starting generation...', 0)
+
+        # Handle VAE decoder switching
+        if use_taehv != current_use_taehv:
+            emit_progress('Switching VAE decoder...', 2)
+            print(f"🔄 Switching VAE decoder to {'TAEHV' if use_taehv else 'default VAE'}")
+            current_vae_decoder = initialize_vae_decoder(use_taehv=use_taehv)
+            # Update pipeline with new VAE decoder
+            pipeline.vae = current_vae_decoder
+
+        # Handle FP8 quantization
+        if enable_fp8 and not fp8_applied:
+            emit_progress('Applying FP8 quantization...', 3)
+            print("🔧 Applying FP8 quantization to transformer")
+            from torchao.quantization.quant_api import quantize_, Float8DynamicActivationFloat8WeightConfig, PerTensor
+            quantize_(transformer, Float8DynamicActivationFloat8WeightConfig(granularity=PerTensor()))
+            fp8_applied = True
+
+        # Text encoding
+        emit_progress('Encoding text prompt...', 8)
+        conditional_dict = text_encoder(text_prompts=[prompt])
+        for key, value in conditional_dict.items():
+            conditional_dict[key] = value.to(dtype=torch.float16)
+        if low_memory:
+            gpu_memory_preservation = get_cuda_free_memory_gb(gpu) + 5
+            move_model_to_device_with_memory_preservation(
+                text_encoder,target_device=gpu, preserved_memory_gb=gpu_memory_preservation)
+
+        # Handle torch.compile if enabled
+        torch_compile_applied = enable_torch_compile
+        if enable_torch_compile and not models_compiled:
+            # Compile transformer and decoder
+            transformer.compile(mode="max-autotune-no-cudagraphs")
+            if not current_use_taehv and not low_memory and not args.trt:
+                current_vae_decoder.compile(mode="max-autotune-no-cudagraphs")
+
+        # Initialize generation
+        emit_progress('Initializing generation...', 12)
+
+        rnd = torch.Generator(gpu).manual_seed(seed)
+        # all_latents = torch.zeros([1, 21, 16, 60, 104], device=gpu, dtype=torch.bfloat16)
+
+        pipeline._initialize_kv_cache(batch_size=1, dtype=torch.float16, device=gpu)
+        pipeline._initialize_crossattn_cache(batch_size=1, dtype=torch.float16, device=gpu)
+
+        noise = torch.randn([1, 21, 16, 60, 104], device=gpu, dtype=torch.float16, generator=rnd)
+
+        # Generation parameters
+        num_blocks = 7
+        current_start_frame = 0
+        num_input_frames = 0
+        all_num_frames = [pipeline.num_frame_per_block] * num_blocks
+        if current_use_taehv:
+            vae_cache = None
+        else:
+            vae_cache = ZERO_VAE_CACHE
+            for i in range(len(vae_cache)):
+                vae_cache[i] = vae_cache[i].to(device=gpu, dtype=torch.float16)
+
+        total_frames_sent = 0
+        generation_start_time = time.time()
+
+        emit_progress('Generating frames... (frontend handles timing)', 15)
+
+        for idx, current_num_frames in enumerate(all_num_frames):
+            if not generation_active or stop_event.is_set():
+                break
+
+            progress = int(((idx + 1) / len(all_num_frames)) * 80) + 15
+
+            # Special message for first block with torch.compile
+            if idx == 0 and torch_compile_applied and not models_compiled:
+                emit_progress(
+                    f'Processing block 1/{len(all_num_frames)} - Compiling models (may take 5-10 minutes)...', progress)
+                print(f"🔥 Processing block {idx+1}/{len(all_num_frames)}")
+                models_compiled = True
+            else:
+                emit_progress(f'Processing block {idx+1}/{len(all_num_frames)}...', progress)
+                print(f"🔄 Processing block {idx+1}/{len(all_num_frames)}")
+
+            block_start_time = time.time()
+
+            noisy_input = noise[:, current_start_frame -
+                                num_input_frames:current_start_frame + current_num_frames - num_input_frames]
+
+            # Denoising loop
+            denoising_start = time.time()
+            for index, current_timestep in enumerate(pipeline.denoising_step_list):
+                if not generation_active or stop_event.is_set():
+                    break
+
+                timestep = torch.ones([1, current_num_frames], device=noise.device,
+                                      dtype=torch.int64) * current_timestep
+
+                if index < len(pipeline.denoising_step_list) - 1:
+                    _, denoised_pred = transformer(
+                        noisy_image_or_video=noisy_input,
+                        conditional_dict=conditional_dict,
+                        timestep=timestep,
+                        kv_cache=pipeline.kv_cache1,
+                        crossattn_cache=pipeline.crossattn_cache,
+                        current_start=current_start_frame * pipeline.frame_seq_length
+                    )
+                    next_timestep = pipeline.denoising_step_list[index + 1]
+                    noisy_input = pipeline.scheduler.add_noise(
+                        denoised_pred.flatten(0, 1),
+                        torch.randn_like(denoised_pred.flatten(0, 1)),
+                        next_timestep * torch.ones([1 * current_num_frames], device=noise.device, dtype=torch.long)
+                    ).unflatten(0, denoised_pred.shape[:2])
+                else:
+                    _, denoised_pred = transformer(
+                        noisy_image_or_video=noisy_input,
+                        conditional_dict=conditional_dict,
+                        timestep=timestep,
+                        kv_cache=pipeline.kv_cache1,
+                        crossattn_cache=pipeline.crossattn_cache,
+                        current_start=current_start_frame * pipeline.frame_seq_length
+                    )
+
+            if not generation_active or stop_event.is_set():
+                break
+
+            denoising_time = time.time() - denoising_start
+            print(f"⚡ Block {idx+1} denoising completed in {denoising_time:.2f}s")
+
+            # Record output
+            # all_latents[:, current_start_frame:current_start_frame + current_num_frames] = denoised_pred
+
+            # Update KV cache for next block
+            if idx != len(all_num_frames) - 1:
+                transformer(
+                    noisy_image_or_video=denoised_pred,
+                    conditional_dict=conditional_dict,
+                    timestep=torch.zeros_like(timestep),
+                    kv_cache=pipeline.kv_cache1,
+                    crossattn_cache=pipeline.crossattn_cache,
+                    current_start=current_start_frame * pipeline.frame_seq_length,
+                )
+
+            # Decode to pixels and send frames immediately
+            print(f"🎨 Decoding block {idx+1} to pixels...")
+            decode_start = time.time()
+            if args.trt:
+                all_current_pixels = []
+                for i in range(denoised_pred.shape[1]):
+                    is_first_frame = torch.tensor(1.0).cuda().half() if idx == 0 and i == 0 else \
+                        torch.tensor(0.0).cuda().half()
+                    outputs = vae_decoder.forward(denoised_pred[:, i:i + 1, :, :, :].half(), is_first_frame, *vae_cache)
+                    # outputs = vae_decoder.forward(denoised_pred.float(), *vae_cache)
+                    current_pixels, vae_cache = outputs[0], outputs[1:]
+                    print(current_pixels.max(), current_pixels.min())
+                    all_current_pixels.append(current_pixels.clone())
+                pixels = torch.cat(all_current_pixels, dim=1)
+                if idx == 0:
+                    pixels = pixels[:, 3:, :, :, :]  # Skip first 3 frames of first block
+            else:
+                if current_use_taehv:
+                    if vae_cache is None:
+                        vae_cache = denoised_pred
+                    else:
+                        denoised_pred = torch.cat([vae_cache, denoised_pred], dim=1)
+                        vae_cache = denoised_pred[:, -3:, :, :, :]
+                    pixels = current_vae_decoder.decode(denoised_pred)
+                    print(f"denoised_pred shape: {denoised_pred.shape}")
+                    print(f"pixels shape: {pixels.shape}")
+                    if idx == 0:
+                        pixels = pixels[:, 3:, :, :, :]  # Skip first 3 frames of first block
+                    else:
+                        pixels = pixels[:, 12:, :, :, :]
+
+                else:
+                    pixels, vae_cache = current_vae_decoder(denoised_pred.half(), *vae_cache)
+                    if idx == 0:
+                        pixels = pixels[:, 3:, :, :, :]  # Skip first 3 frames of first block
+
+            decode_time = time.time() - decode_start
+            print(f"🎨 Block {idx+1} VAE decoding completed in {decode_time:.2f}s")
+
+            # Queue frames for non-blocking sending
+            block_frames = pixels.shape[1]
+            print(f"📡 Queueing {block_frames} frames from block {idx+1} for sending...")
+            queue_start = time.time()
+
+            for frame_idx in range(block_frames):
+                if not generation_active or stop_event.is_set():
+                    break
+
+                frame_tensor = pixels[0, frame_idx].cpu()
+
+                # Queue frame data in non-blocking way
+                frame_send_queue.put((frame_tensor, total_frames_sent, idx, job_id))
+                total_frames_sent += 1
+
+            queue_time = time.time() - queue_start
+            block_time = time.time() - block_start_time
+            print(f"✅ Block {idx+1} completed in {block_time:.2f}s ({block_frames} frames queued in {queue_time:.3f}s)")
+
+            current_start_frame += current_num_frames
+
+        generation_time = time.time() - generation_start_time
+        print(f"🎉 Generation completed in {generation_time:.2f}s! {total_frames_sent} frames queued for sending")
+
+        # Wait for all frames to be sent before completing
+        emit_progress('Waiting for all frames to be sent...', 97)
+        print("⏳ Waiting for all frames to be sent...")
+        frame_send_queue.join()  # Wait for all queued frames to be processed
+        print("✅ All frames sent successfully!")
+
+        generate_mp4_from_images("./images","./videos/"+anim_name+".mp4", frame_rate )
+        # Final progress update
+        emit_progress('Generation complete!', 100)
+
+        try:
+            socketio.emit('generation_complete', {
+                'message': 'Video generation completed!',
+                'total_frames': total_frames_sent,
+                'generation_time': f"{generation_time:.2f}s",
+                'job_id': job_id
+            })
+        except Exception as e:
+            print(f"❌ Failed to emit generation complete: {e}")
+
+    except Exception as e:
+        print(f"❌ Generation failed: {e}")
+        try:
+            socketio.emit('error', {
+                'message': f'Generation failed: {str(e)}',
+                'job_id': job_id
+            })
+        except Exception as e:
+            print(f"❌ Failed to emit error: {e}")
+    finally:
+        generation_active = False
+        stop_event.set()
+
+        # Clean up sender thread
+        try:
+            frame_send_queue.put(None)
+        except Exception as e:
+            print(f"❌ Failed to put None in frame_send_queue: {e}")
+
+
+def generate_mp4_from_images(image_directory, output_video_path, fps=24):
+    """
+    Generate an MP4 video from a directory of images ordered alphabetically.
+
+    :param image_directory: Path to the directory containing images.
+    :param output_video_path: Path where the output MP4 will be saved.
+    :param fps: Frames per second for the output video.
+    """
+    global anim_name
+    # Construct the ffmpeg command
+    cmd = [
+        'ffmpeg',
+        '-framerate', str(fps),
+        '-i', os.path.join(image_directory, anim_name+'/'+anim_name+'_%03d.jpg'),  # Adjust the pattern if necessary
+        '-c:v', 'libx264',
+        '-pix_fmt', 'yuv420p',
+        output_video_path
+    ]
+    try:
+        subprocess.run(cmd, check=True)
+        print(f"Video saved to {output_video_path}")
+    except subprocess.CalledProcessError as e:
+        print(f"An error occurred: {e}")
+
+def calculate_sha256(data):
+    # Convert data to bytes if it's not already
+    if isinstance(data, str):
+        data = data.encode()
+    # Calculate SHA-256 hash
+    sha256_hash = hashlib.sha256(data).hexdigest()
+    return sha256_hash
+
+# Socket.IO event handlers
+@socketio.on('connect')
+def handle_connect():
+    print('Client connected')
+    emit('status', {'message': 'Connected to frontend-buffered demo server'})
+
+
+@socketio.on('disconnect')
+def handle_disconnect():
+    print('Client disconnected')
+
+
+@socketio.on('start_generation')
+def handle_start_generation(data):
+    global generation_active, frame_number, anim_name, frame_rate
+
+    frame_number = 0
+    if generation_active:
+        emit('error', {'message': 'Generation already in progress'})
+        return
+
+    prompt = data.get('prompt', '')
+
+    seed = data.get('seed', -1)
+    if seed==-1:
+        seed = random.randint(0, 2**32)
+
+    # Extract words up to the first punctuation or newline
+    words_up_to_punctuation = re.split(r'[^\w\s]', prompt)[0].strip() if prompt else ''
+    if not words_up_to_punctuation:
+        words_up_to_punctuation = re.split(r'[\n\r]', prompt)[0].strip()
+
+    # Calculate SHA-256 hash of the entire prompt
+    sha256_hash = calculate_sha256(prompt)
+
+    # Create anim_name with the extracted words and first 10 characters of the hash
+    anim_name = f"{words_up_to_punctuation[:20]}_{str(seed)}_{sha256_hash[:10]}"
+
+    generation_active = True
+    generation_start_time = time.time()
+    enable_torch_compile = data.get('enable_torch_compile', False)
+    enable_fp8 = data.get('enable_fp8', False)
+    use_taehv = data.get('use_taehv', False)
+    frame_rate = data.get('fps', 6)
+
+    if not prompt:
+        emit('error', {'message': 'Prompt is required'})
+        return
+
+    # Start generation in background thread
+    socketio.start_background_task(generate_video_stream, prompt, seed,
+                                   enable_torch_compile, enable_fp8, use_taehv)
+    emit('status', {'message': 'Generation started - frames will be sent immediately'})
+
+
+@socketio.on('stop_generation')
+def handle_stop_generation():
+    global generation_active, stop_event, frame_send_queue
+    generation_active = False
+    stop_event.set()
+
+    # Signal sender thread to stop (will be processed after current frames)
+    try:
+        frame_send_queue.put(None)
+    except Exception as e:
+        print(f"❌ Failed to put None in frame_send_queue: {e}")
+
+    emit('status', {'message': 'Generation stopped'})
+
+# Web routes
+
+
+@app.route('/')
+def index():
+    return render_template('demo.html')
+
+
+@app.route('/api/status')
+def api_status():
+    return jsonify({
+        'generation_active': generation_active,
+        'free_vram_gb': get_cuda_free_memory_gb(gpu),
+        'fp8_applied': fp8_applied,
+        'torch_compile_applied': torch_compile_applied,
+        'current_use_taehv': current_use_taehv
+    })
+
+
+if __name__ == '__main__':
+    print(f"🚀 Starting demo on http://{args.host}:{args.port}")
+    socketio.run(app, host=args.host, port=args.port, debug=False)

model/__init__.py

@@ -0,0 +1,36 @@
+"""Public model exports loaded on demand.
+
+Lazy imports preserve the existing ``from model import DMD`` API while avoiding
+the model/pipeline cycle introduced by the standalone Predictor module.
+"""
+
+from __future__ import annotations
+
+from importlib import import_module
+
+
+_EXPORTS = {
+    "CausalDiffusion": ("model.diffusion", "CausalDiffusion"),
+    "CausVid": ("model.causvid", "CausVid"),
+    "DMD": ("model.dmd", "DMD"),
+    "GAN": ("model.gan", "GAN"),
+    "SiD": ("model.sid", "SiD"),
+    "ODERegression": ("model.ode_regression", "ODERegression"),
+    "SelfForcingPredictorV4": (
+        "model.predictor_v4",
+        "SelfForcingPredictorV4",
+    ),
+    "TripleFeatureFusion": ("model.predictor_v4", "TripleFeatureFusion"),
+    "WanPredictorV4Config": ("model.predictor_v4", "WanPredictorV4Config"),
+}
+
+__all__ = list(_EXPORTS)
+
+
+def __getattr__(name: str):
+    if name not in _EXPORTS:
+        raise AttributeError(name)
+    module_name, attribute = _EXPORTS[name]
+    value = getattr(import_module(module_name), attribute)
+    globals()[name] = value
+    return value

model/base.py

@@ -0,0 +1,222 @@
+from typing import Tuple
+from einops import rearrange
+from torch import nn
+import torch.distributed as dist
+import torch
+
+from pipeline import SelfForcingTrainingPipeline
+from utils.loss import get_denoising_loss
+from utils.wan_wrapper import WanDiffusionWrapper, WanTextEncoder, WanVAEWrapper
+
+
+class BaseModel(nn.Module):
+    def __init__(self, args, device):
+        super().__init__()
+        self._initialize_models(args, device)
+
+        self.device = device
+        self.args = args
+        self.dtype = torch.bfloat16 if args.mixed_precision else torch.float32
+        if hasattr(args, "denoising_step_list"):
+            self.denoising_step_list = torch.tensor(args.denoising_step_list, dtype=torch.long)
+            if args.warp_denoising_step:
+                timesteps = torch.cat((self.scheduler.timesteps.cpu(), torch.tensor([0], dtype=torch.float32)))
+                self.denoising_step_list = timesteps[1000 - self.denoising_step_list]
+
+    def _initialize_models(self, args, device):
+        self.real_model_name = getattr(args, "real_name", "Wan2.1-T2V-1.3B")
+        self.fake_model_name = getattr(args, "fake_name", "Wan2.1-T2V-1.3B")
+
+        self.generator = WanDiffusionWrapper(**getattr(args, "model_kwargs", {}), is_causal=True)
+        self.generator.model.requires_grad_(True)
+
+        self.real_score = WanDiffusionWrapper(model_name=self.real_model_name, is_causal=False)
+        self.real_score.model.requires_grad_(False)
+
+        self.fake_score = WanDiffusionWrapper(model_name=self.fake_model_name, is_causal=False)
+        self.fake_score.model.requires_grad_(True)
+
+        self.text_encoder = WanTextEncoder()
+        self.text_encoder.requires_grad_(False)
+
+        self.vae = WanVAEWrapper()
+        self.vae.requires_grad_(False)
+
+        self.scheduler = self.generator.get_scheduler()
+        self.scheduler.timesteps = self.scheduler.timesteps.to(device)
+
+    def _get_timestep(
+            self,
+            min_timestep: int,
+            max_timestep: int,
+            batch_size: int,
+            num_frame: int,
+            num_frame_per_block: int,
+            uniform_timestep: bool = False
+    ) -> torch.Tensor:
+        """
+        Randomly generate a timestep tensor based on the generator's task type. It uniformly samples a timestep
+        from the range [min_timestep, max_timestep], and returns a tensor of shape [batch_size, num_frame].
+        - If uniform_timestep, it will use the same timestep for all frames.
+        - If not uniform_timestep, it will use a different timestep for each block.
+        """
+        if uniform_timestep:
+            timestep = torch.randint(
+                min_timestep,
+                max_timestep,
+                [batch_size, 1],
+                device=self.device,
+                dtype=torch.long
+            ).repeat(1, num_frame)
+            return timestep
+        else:
+            timestep = torch.randint(
+                min_timestep,
+                max_timestep,
+                [batch_size, num_frame],
+                device=self.device,
+                dtype=torch.long
+            )
+            # make the noise level the same within every block
+            if self.independent_first_frame:
+                # the first frame is always kept the same
+                timestep_from_second = timestep[:, 1:]
+                timestep_from_second = timestep_from_second.reshape(
+                    timestep_from_second.shape[0], -1, num_frame_per_block)
+                timestep_from_second[:, :, 1:] = timestep_from_second[:, :, 0:1]
+                timestep_from_second = timestep_from_second.reshape(
+                    timestep_from_second.shape[0], -1)
+                timestep = torch.cat([timestep[:, 0:1], timestep_from_second], dim=1)
+            else:
+                timestep = timestep.reshape(
+                    timestep.shape[0], -1, num_frame_per_block)
+                timestep[:, :, 1:] = timestep[:, :, 0:1]
+                timestep = timestep.reshape(timestep.shape[0], -1)
+            return timestep
+
+
+class SelfForcingModel(BaseModel):
+    def __init__(self, args, device):
+        super().__init__(args, device)
+        self.denoising_loss_func = get_denoising_loss(args.denoising_loss_type)()
+
+    def _run_generator(
+        self,
+        image_or_video_shape,
+        conditional_dict: dict,
+        initial_latent: torch.tensor = None
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Optionally simulate the generator's input from noise using backward simulation
+        and then run the generator for one-step.
+        Input:
+            - image_or_video_shape: a list containing the shape of the image or video [B, F, C, H, W].
+            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
+            - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
+            - clean_latent: a tensor containing the clean latents [B, F, C, H, W]. Need to be passed when no backward simulation is used.
+            - initial_latent: a tensor containing the initial latents [B, F, C, H, W].
+        Output:
+            - pred_image: a tensor with shape [B, F, C, H, W].
+            - denoised_timestep: an integer
+        """
+        # Step 1: Sample noise and backward simulate the generator's input
+        assert getattr(self.args, "backward_simulation", True), "Backward simulation needs to be enabled"
+        if initial_latent is not None:
+            conditional_dict["initial_latent"] = initial_latent
+        if self.args.i2v:
+            noise_shape = [image_or_video_shape[0], image_or_video_shape[1] - 1, *image_or_video_shape[2:]]
+        else:
+            noise_shape = image_or_video_shape.copy()
+
+        # During training, the number of generated frames should be uniformly sampled from
+        # [21, self.num_training_frames], but still being a multiple of self.num_frame_per_block
+        min_num_frames = 20 if self.args.independent_first_frame else 21
+        max_num_frames = self.num_training_frames - 1 if self.args.independent_first_frame else self.num_training_frames
+        assert max_num_frames % self.num_frame_per_block == 0
+        assert min_num_frames % self.num_frame_per_block == 0
+        max_num_blocks = max_num_frames // self.num_frame_per_block
+        min_num_blocks = min_num_frames // self.num_frame_per_block
+        num_generated_blocks = torch.randint(min_num_blocks, max_num_blocks + 1, (1,), device=self.device)
+        dist.broadcast(num_generated_blocks, src=0)
+        num_generated_blocks = num_generated_blocks.item()
+        num_generated_frames = num_generated_blocks * self.num_frame_per_block
+        if self.args.independent_first_frame and initial_latent is None:
+            num_generated_frames += 1
+            min_num_frames += 1
+        # Sync num_generated_frames across all processes
+        noise_shape[1] = num_generated_frames
+
+        pred_image_or_video, denoised_timestep_from, denoised_timestep_to = self._consistency_backward_simulation(
+            noise=torch.randn(noise_shape,
+                              device=self.device, dtype=self.dtype),
+            **conditional_dict,
+        )
+        # Slice last 21 frames
+        if pred_image_or_video.shape[1] > 21:
+            with torch.no_grad():
+                # Reencode to get image latent
+                latent_to_decode = pred_image_or_video[:, :-20, ...]
+                # Deccode to video
+                pixels = self.vae.decode_to_pixel(latent_to_decode)
+                frame = pixels[:, -1:, ...].to(self.dtype)
+                frame = rearrange(frame, "b t c h w -> b c t h w")
+                # Encode frame to get image latent
+                image_latent = self.vae.encode_to_latent(frame).to(self.dtype)
+            pred_image_or_video_last_21 = torch.cat([image_latent, pred_image_or_video[:, -20:, ...]], dim=1)
+        else:
+            pred_image_or_video_last_21 = pred_image_or_video
+
+        if num_generated_frames != min_num_frames:
+            # Currently, we do not use gradient for the first chunk, since it contains image latents
+            gradient_mask = torch.ones_like(pred_image_or_video_last_21, dtype=torch.bool)
+            if self.args.independent_first_frame:
+                gradient_mask[:, :1] = False
+            else:
+                gradient_mask[:, :self.num_frame_per_block] = False
+        else:
+            gradient_mask = None
+
+        pred_image_or_video_last_21 = pred_image_or_video_last_21.to(self.dtype)
+        return pred_image_or_video_last_21, gradient_mask, denoised_timestep_from, denoised_timestep_to
+
+    def _consistency_backward_simulation(
+        self,
+        noise: torch.Tensor,
+        **conditional_dict: dict
+    ) -> torch.Tensor:
+        """
+        Simulate the generator's input from noise to avoid training/inference mismatch.
+        See Sec 4.5 of the DMD2 paper (https://arxiv.org/abs/2405.14867) for details.
+        Here we use the consistency sampler (https://arxiv.org/abs/2303.01469)
+        Input:
+            - noise: a tensor sampled from N(0, 1) with shape [B, F, C, H, W] where the number of frame is 1 for images.
+            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
+        Output:
+            - output: a tensor with shape [B, T, F, C, H, W].
+            T is the total number of timesteps. output[0] is a pure noise and output[i] and i>0
+            represents the x0 prediction at each timestep.
+        """
+        if self.inference_pipeline is None:
+            self._initialize_inference_pipeline()
+
+        return self.inference_pipeline.inference_with_trajectory(
+            noise=noise, **conditional_dict
+        )
+
+    def _initialize_inference_pipeline(self):
+        """
+        Lazy initialize the inference pipeline during the first backward simulation run.
+        Here we encapsulate the inference code with a model-dependent outside function.
+        We pass our FSDP-wrapped modules into the pipeline to save memory.
+        """
+        self.inference_pipeline = SelfForcingTrainingPipeline(
+            denoising_step_list=self.denoising_step_list,
+            scheduler=self.scheduler,
+            generator=self.generator,
+            num_frame_per_block=self.num_frame_per_block,
+            independent_first_frame=self.args.independent_first_frame,
+            same_step_across_blocks=self.args.same_step_across_blocks,
+            last_step_only=self.args.last_step_only,
+            num_max_frames=self.num_training_frames,
+            context_noise=self.args.context_noise
+        )

model/causvid.py

@@ -0,0 +1,391 @@
+import torch.nn.functional as F
+from typing import Tuple
+import torch
+
+from model.base import BaseModel
+
+
+class CausVid(BaseModel):
+    def __init__(self, args, device):
+        """
+        Initialize the DMD (Distribution Matching Distillation) module.
+        This class is self-contained and compute generator and fake score losses
+        in the forward pass.
+        """
+        super().__init__(args, device)
+        self.num_frame_per_block = getattr(args, "num_frame_per_block", 1)
+        self.num_training_frames = getattr(args, "num_training_frames", 21)
+
+        if self.num_frame_per_block > 1:
+            self.generator.model.num_frame_per_block = self.num_frame_per_block
+
+        self.independent_first_frame = getattr(args, "independent_first_frame", False)
+        if self.independent_first_frame:
+            self.generator.model.independent_first_frame = True
+        if args.gradient_checkpointing:
+            self.generator.enable_gradient_checkpointing()
+            self.fake_score.enable_gradient_checkpointing()
+
+        # Step 2: Initialize all dmd hyperparameters
+        self.num_train_timestep = args.num_train_timestep
+        self.min_step = int(0.02 * self.num_train_timestep)
+        self.max_step = int(0.98 * self.num_train_timestep)
+        if hasattr(args, "real_guidance_scale"):
+            self.real_guidance_scale = args.real_guidance_scale
+            self.fake_guidance_scale = args.fake_guidance_scale
+        else:
+            self.real_guidance_scale = args.guidance_scale
+            self.fake_guidance_scale = 0.0
+        self.timestep_shift = getattr(args, "timestep_shift", 1.0)
+        self.teacher_forcing = getattr(args, "teacher_forcing", False)
+
+        if getattr(self.scheduler, "alphas_cumprod", None) is not None:
+            self.scheduler.alphas_cumprod = self.scheduler.alphas_cumprod.to(device)
+        else:
+            self.scheduler.alphas_cumprod = None
+
+    def _compute_kl_grad(
+        self, noisy_image_or_video: torch.Tensor,
+        estimated_clean_image_or_video: torch.Tensor,
+        timestep: torch.Tensor,
+        conditional_dict: dict, unconditional_dict: dict,
+        normalization: bool = True
+    ) -> Tuple[torch.Tensor, dict]:
+        """
+        Compute the KL grad (eq 7 in https://arxiv.org/abs/2311.18828).
+        Input:
+            - noisy_image_or_video: a tensor with shape [B, F, C, H, W] where the number of frame is 1 for images.
+            - estimated_clean_image_or_video: a tensor with shape [B, F, C, H, W] representing the estimated clean image or video.
+            - timestep: a tensor with shape [B, F] containing the randomly generated timestep.
+            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
+            - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
+            - normalization: a boolean indicating whether to normalize the gradient.
+        Output:
+            - kl_grad: a tensor representing the KL grad.
+            - kl_log_dict: a dictionary containing the intermediate tensors for logging.
+        """
+        # Step 1: Compute the fake score
+        _, pred_fake_image_cond = self.fake_score(
+            noisy_image_or_video=noisy_image_or_video,
+            conditional_dict=conditional_dict,
+            timestep=timestep
+        )
+
+        if self.fake_guidance_scale != 0.0:
+            _, pred_fake_image_uncond = self.fake_score(
+                noisy_image_or_video=noisy_image_or_video,
+                conditional_dict=unconditional_dict,
+                timestep=timestep
+            )
+            pred_fake_image = pred_fake_image_cond + (
+                pred_fake_image_cond - pred_fake_image_uncond
+            ) * self.fake_guidance_scale
+        else:
+            pred_fake_image = pred_fake_image_cond
+
+        # Step 2: Compute the real score
+        # We compute the conditional and unconditional prediction
+        # and add them together to achieve cfg (https://arxiv.org/abs/2207.12598)
+        _, pred_real_image_cond = self.real_score(
+            noisy_image_or_video=noisy_image_or_video,
+            conditional_dict=conditional_dict,
+            timestep=timestep
+        )
+
+        _, pred_real_image_uncond = self.real_score(
+            noisy_image_or_video=noisy_image_or_video,
+            conditional_dict=unconditional_dict,
+            timestep=timestep
+        )
+
+        pred_real_image = pred_real_image_cond + (
+            pred_real_image_cond - pred_real_image_uncond
+        ) * self.real_guidance_scale
+
+        # Step 3: Compute the DMD gradient (DMD paper eq. 7).
+        grad = (pred_fake_image - pred_real_image)
+
+        # TODO: Change the normalizer for causal teacher
+        if normalization:
+            # Step 4: Gradient normalization (DMD paper eq. 8).
+            p_real = (estimated_clean_image_or_video - pred_real_image)
+            normalizer = torch.abs(p_real).mean(dim=[1, 2, 3, 4], keepdim=True)
+            grad = grad / normalizer
+        grad = torch.nan_to_num(grad)
+
+        return grad, {
+            "dmdtrain_gradient_norm": torch.mean(torch.abs(grad)).detach(),
+            "timestep": timestep.detach()
+        }
+
+    def compute_distribution_matching_loss(
+        self,
+        image_or_video: torch.Tensor,
+        conditional_dict: dict,
+        unconditional_dict: dict,
+        gradient_mask: torch.Tensor = None,
+    ) -> Tuple[torch.Tensor, dict]:
+        """
+        Compute the DMD loss (eq 7 in https://arxiv.org/abs/2311.18828).
+        Input:
+            - image_or_video: a tensor with shape [B, F, C, H, W] where the number of frame is 1 for images.
+            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
+            - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
+            - gradient_mask: a boolean tensor with the same shape as image_or_video indicating which pixels to compute loss .
+        Output:
+            - dmd_loss: a scalar tensor representing the DMD loss.
+            - dmd_log_dict: a dictionary containing the intermediate tensors for logging.
+        """
+        original_latent = image_or_video
+
+        batch_size, num_frame = image_or_video.shape[:2]
+
+        with torch.no_grad():
+            # Step 1: Randomly sample timestep based on the given schedule and corresponding noise
+            timestep = self._get_timestep(
+                0,
+                self.num_train_timestep,
+                batch_size,
+                num_frame,
+                self.num_frame_per_block,
+                uniform_timestep=True
+            )
+
+            if self.timestep_shift > 1:
+                timestep = self.timestep_shift * \
+                    (timestep / 1000) / \
+                    (1 + (self.timestep_shift - 1) * (timestep / 1000)) * 1000
+            timestep = timestep.clamp(self.min_step, self.max_step)
+
+            noise = torch.randn_like(image_or_video)
+            noisy_latent = self.scheduler.add_noise(
+                image_or_video.flatten(0, 1),
+                noise.flatten(0, 1),
+                timestep.flatten(0, 1)
+            ).detach().unflatten(0, (batch_size, num_frame))
+
+            # Step 2: Compute the KL grad
+            grad, dmd_log_dict = self._compute_kl_grad(
+                noisy_image_or_video=noisy_latent,
+                estimated_clean_image_or_video=original_latent,
+                timestep=timestep,
+                conditional_dict=conditional_dict,
+                unconditional_dict=unconditional_dict
+            )
+
+        if gradient_mask is not None:
+            dmd_loss = 0.5 * F.mse_loss(original_latent.double(
+            )[gradient_mask], (original_latent.double() - grad.double()).detach()[gradient_mask], reduction="mean")
+        else:
+            dmd_loss = 0.5 * F.mse_loss(original_latent.double(
+            ), (original_latent.double() - grad.double()).detach(), reduction="mean")
+        return dmd_loss, dmd_log_dict
+
+    def _run_generator(
+        self,
+        image_or_video_shape,
+        conditional_dict: dict,
+        clean_latent: torch.tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Optionally simulate the generator's input from noise using backward simulation
+        and then run the generator for one-step.
+        Input:
+            - image_or_video_shape: a list containing the shape of the image or video [B, F, C, H, W].
+            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
+            - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
+            - clean_latent: a tensor containing the clean latents [B, F, C, H, W]. Need to be passed when no backward simulation is used.
+            - initial_latent: a tensor containing the initial latents [B, F, C, H, W].
+        Output:
+            - pred_image: a tensor with shape [B, F, C, H, W].
+        """
+        simulated_noisy_input = []
+        for timestep in self.denoising_step_list:
+            noise = torch.randn(
+                image_or_video_shape, device=self.device, dtype=self.dtype)
+
+            noisy_timestep = timestep * torch.ones(
+                image_or_video_shape[:2], device=self.device, dtype=torch.long)
+
+            if timestep != 0:
+                noisy_image = self.scheduler.add_noise(
+                    clean_latent.flatten(0, 1),
+                    noise.flatten(0, 1),
+                    noisy_timestep.flatten(0, 1)
+                ).unflatten(0, image_or_video_shape[:2])
+            else:
+                noisy_image = clean_latent
+
+            simulated_noisy_input.append(noisy_image)
+
+        simulated_noisy_input = torch.stack(simulated_noisy_input, dim=1)
+
+        # Step 2: Randomly sample a timestep and pick the corresponding input
+        index = self._get_timestep(
+            0,
+            len(self.denoising_step_list),
+            image_or_video_shape[0],
+            image_or_video_shape[1],
+            self.num_frame_per_block,
+            uniform_timestep=False
+        )
+
+        # select the corresponding timestep's noisy input from the stacked tensor [B, T, F, C, H, W]
+        noisy_input = torch.gather(
+            simulated_noisy_input, dim=1,
+            index=index.reshape(index.shape[0], 1, index.shape[1], 1, 1, 1).expand(
+                -1, -1, -1, *image_or_video_shape[2:]).to(self.device)
+        ).squeeze(1)
+
+        timestep = self.denoising_step_list[index].to(self.device)
+
+        _, pred_image_or_video = self.generator(
+            noisy_image_or_video=noisy_input,
+            conditional_dict=conditional_dict,
+            timestep=timestep,
+            clean_x=clean_latent if self.teacher_forcing else None,
+        )
+
+        gradient_mask = None  # timestep != 0
+
+        pred_image_or_video = pred_image_or_video.type_as(noisy_input)
+
+        return pred_image_or_video, gradient_mask
+
+    def generator_loss(
+        self,
+        image_or_video_shape,
+        conditional_dict: dict,
+        unconditional_dict: dict,
+        clean_latent: torch.Tensor,
+        initial_latent: torch.Tensor = None
+    ) -> Tuple[torch.Tensor, dict]:
+        """
+        Generate image/videos from noise and compute the DMD loss.
+        The noisy input to the generator is backward simulated.
+        This removes the need of any datasets during distillation.
+        See Sec 4.5 of the DMD2 paper (https://arxiv.org/abs/2405.14867) for details.
+        Input:
+            - image_or_video_shape: a list containing the shape of the image or video [B, F, C, H, W].
+            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
+            - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
+            - clean_latent: a tensor containing the clean latents [B, F, C, H, W]. Need to be passed when no backward simulation is used.
+        Output:
+            - loss: a scalar tensor representing the generator loss.
+            - generator_log_dict: a dictionary containing the intermediate tensors for logging.
+        """
+        # Step 1: Run generator on backward simulated noisy input
+        pred_image, gradient_mask = self._run_generator(
+            image_or_video_shape=image_or_video_shape,
+            conditional_dict=conditional_dict,
+            clean_latent=clean_latent
+        )
+
+        # Step 2: Compute the DMD loss
+        dmd_loss, dmd_log_dict = self.compute_distribution_matching_loss(
+            image_or_video=pred_image,
+            conditional_dict=conditional_dict,
+            unconditional_dict=unconditional_dict,
+            gradient_mask=gradient_mask
+        )
+
+        # Step 3: TODO: Implement the GAN loss
+
+        return dmd_loss, dmd_log_dict
+
+    def critic_loss(
+        self,
+        image_or_video_shape,
+        conditional_dict: dict,
+        unconditional_dict: dict,
+        clean_latent: torch.Tensor,
+        initial_latent: torch.Tensor = None
+    ) -> Tuple[torch.Tensor, dict]:
+        """
+        Generate image/videos from noise and train the critic with generated samples.
+        The noisy input to the generator is backward simulated.
+        This removes the need of any datasets during distillation.
+        See Sec 4.5 of the DMD2 paper (https://arxiv.org/abs/2405.14867) for details.
+        Input:
+            - image_or_video_shape: a list containing the shape of the image or video [B, F, C, H, W].
+            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
+            - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
+            - clean_latent: a tensor containing the clean latents [B, F, C, H, W]. Need to be passed when no backward simulation is used.
+        Output:
+            - loss: a scalar tensor representing the generator loss.
+            - critic_log_dict: a dictionary containing the intermediate tensors for logging.
+        """
+
+        # Step 1: Run generator on backward simulated noisy input
+        with torch.no_grad():
+            generated_image, _ = self._run_generator(
+                image_or_video_shape=image_or_video_shape,
+                conditional_dict=conditional_dict,
+                clean_latent=clean_latent
+            )
+
+        # Step 2: Compute the fake prediction
+        critic_timestep = self._get_timestep(
+            0,
+            self.num_train_timestep,
+            image_or_video_shape[0],
+            image_or_video_shape[1],
+            self.num_frame_per_block,
+            uniform_timestep=True
+        )
+
+        if self.timestep_shift > 1:
+            critic_timestep = self.timestep_shift * \
+                (critic_timestep / 1000) / (1 + (self.timestep_shift - 1) * (critic_timestep / 1000)) * 1000
+
+        critic_timestep = critic_timestep.clamp(self.min_step, self.max_step)
+
+        critic_noise = torch.randn_like(generated_image)
+        noisy_generated_image = self.scheduler.add_noise(
+            generated_image.flatten(0, 1),
+            critic_noise.flatten(0, 1),
+            critic_timestep.flatten(0, 1)
+        ).unflatten(0, image_or_video_shape[:2])
+
+        _, pred_fake_image = self.fake_score(
+            noisy_image_or_video=noisy_generated_image,
+            conditional_dict=conditional_dict,
+            timestep=critic_timestep
+        )
+
+        # Step 3: Compute the denoising loss for the fake critic
+        if self.args.denoising_loss_type == "flow":
+            from utils.wan_wrapper import WanDiffusionWrapper
+            flow_pred = WanDiffusionWrapper._convert_x0_to_flow_pred(
+                scheduler=self.scheduler,
+                x0_pred=pred_fake_image.flatten(0, 1),
+                xt=noisy_generated_image.flatten(0, 1),
+                timestep=critic_timestep.flatten(0, 1)
+            )
+            pred_fake_noise = None
+        else:
+            flow_pred = None
+            pred_fake_noise = self.scheduler.convert_x0_to_noise(
+                x0=pred_fake_image.flatten(0, 1),
+                xt=noisy_generated_image.flatten(0, 1),
+                timestep=critic_timestep.flatten(0, 1)
+            ).unflatten(0, image_or_video_shape[:2])
+
+        denoising_loss = self.denoising_loss_func(
+            x=generated_image.flatten(0, 1),
+            x_pred=pred_fake_image.flatten(0, 1),
+            noise=critic_noise.flatten(0, 1),
+            noise_pred=pred_fake_noise,
+            alphas_cumprod=self.scheduler.alphas_cumprod,
+            timestep=critic_timestep.flatten(0, 1),
+            flow_pred=flow_pred
+        )
+
+        # Step 4: TODO: Compute the GAN loss
+
+        # Step 5: Debugging Log
+        critic_log_dict = {
+            "critic_timestep": critic_timestep.detach()
+        }
+
+        return denoising_loss, critic_log_dict

model/diffusion.py

@@ -0,0 +1,125 @@
+from typing import Tuple
+import torch
+
+from model.base import BaseModel
+from utils.wan_wrapper import WanDiffusionWrapper, WanTextEncoder, WanVAEWrapper
+
+
+class CausalDiffusion(BaseModel):
+    def __init__(self, args, device):
+        """
+        Initialize the Diffusion loss module.
+        """
+        super().__init__(args, device)
+        self.num_frame_per_block = getattr(args, "num_frame_per_block", 1)
+        if self.num_frame_per_block > 1:
+            self.generator.model.num_frame_per_block = self.num_frame_per_block
+        self.independent_first_frame = getattr(args, "independent_first_frame", False)
+        if self.independent_first_frame:
+            self.generator.model.independent_first_frame = True
+
+        if args.gradient_checkpointing:
+            self.generator.enable_gradient_checkpointing()
+
+        # Step 2: Initialize all hyperparameters
+        self.num_train_timestep = args.num_train_timestep
+        self.min_step = int(0.02 * self.num_train_timestep)
+        self.max_step = int(0.98 * self.num_train_timestep)
+        self.guidance_scale = args.guidance_scale
+        self.timestep_shift = getattr(args, "timestep_shift", 1.0)
+        self.teacher_forcing = getattr(args, "teacher_forcing", False)
+        # Noise augmentation in teacher forcing, we add small noise to clean context latents
+        self.noise_augmentation_max_timestep = getattr(args, "noise_augmentation_max_timestep", 0)
+
+    def _initialize_models(self, args):
+        self.generator = WanDiffusionWrapper(**getattr(args, "model_kwargs", {}), is_causal=True)
+        self.generator.model.requires_grad_(True)
+
+        self.text_encoder = WanTextEncoder()
+        self.text_encoder.requires_grad_(False)
+
+        self.vae = WanVAEWrapper()
+        self.vae.requires_grad_(False)
+
+    def generator_loss(
+        self,
+        image_or_video_shape,
+        conditional_dict: dict,
+        unconditional_dict: dict,
+        clean_latent: torch.Tensor,
+        initial_latent: torch.Tensor = None
+    ) -> Tuple[torch.Tensor, dict]:
+        """
+        Generate image/videos from noise and compute the DMD loss.
+        The noisy input to the generator is backward simulated.
+        This removes the need of any datasets during distillation.
+        See Sec 4.5 of the DMD2 paper (https://arxiv.org/abs/2405.14867) for details.
+        Input:
+            - image_or_video_shape: a list containing the shape of the image or video [B, F, C, H, W].
+            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
+            - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
+            - clean_latent: a tensor containing the clean latents [B, F, C, H, W]. Need to be passed when no backward simulation is used.
+        Output:
+            - loss: a scalar tensor representing the generator loss.
+            - generator_log_dict: a dictionary containing the intermediate tensors for logging.
+        """
+        noise = torch.randn_like(clean_latent)
+        batch_size, num_frame = image_or_video_shape[:2]
+
+        # Step 2: Randomly sample a timestep and add noise to denoiser inputs
+        index = self._get_timestep(
+            0,
+            self.scheduler.num_train_timesteps,
+            image_or_video_shape[0],
+            image_or_video_shape[1],
+            self.num_frame_per_block,
+            uniform_timestep=False
+        )
+        timestep = self.scheduler.timesteps[index].to(dtype=self.dtype, device=self.device)
+        noisy_latents = self.scheduler.add_noise(
+            clean_latent.flatten(0, 1),
+            noise.flatten(0, 1),
+            timestep.flatten(0, 1)
+        ).unflatten(0, (batch_size, num_frame))
+        training_target = self.scheduler.training_target(clean_latent, noise, timestep)
+
+        # Step 3: Noise augmentation, also add small noise to clean context latents
+        if self.noise_augmentation_max_timestep > 0:
+            index_clean_aug = self._get_timestep(
+                0,
+                self.noise_augmentation_max_timestep,
+                image_or_video_shape[0],
+                image_or_video_shape[1],
+                self.num_frame_per_block,
+                uniform_timestep=False
+            )
+            timestep_clean_aug = self.scheduler.timesteps[index_clean_aug].to(dtype=self.dtype, device=self.device)
+            clean_latent_aug = self.scheduler.add_noise(
+                clean_latent.flatten(0, 1),
+                noise.flatten(0, 1),
+                timestep_clean_aug.flatten(0, 1)
+            ).unflatten(0, (batch_size, num_frame))
+        else:
+            clean_latent_aug = clean_latent
+            timestep_clean_aug = None
+
+        # Compute loss
+        flow_pred, x0_pred = self.generator(
+            noisy_image_or_video=noisy_latents,
+            conditional_dict=conditional_dict,
+            timestep=timestep,
+            clean_x=clean_latent_aug if self.teacher_forcing else None,
+            aug_t=timestep_clean_aug if self.teacher_forcing else None
+        )
+        # loss = torch.nn.functional.mse_loss(flow_pred.float(), training_target.float())
+        loss = torch.nn.functional.mse_loss(
+            flow_pred.float(), training_target.float(), reduction='none'
+        ).mean(dim=(2, 3, 4))
+        loss = loss * self.scheduler.training_weight(timestep).unflatten(0, (batch_size, num_frame))
+        loss = loss.mean()
+
+        log_dict = {
+            "x0": clean_latent.detach(),
+            "x0_pred": x0_pred.detach()
+        }
+        return loss, log_dict

model/dmd.py

@@ -0,0 +1,332 @@
+from pipeline import SelfForcingTrainingPipeline
+import torch.nn.functional as F
+from typing import Optional, Tuple
+import torch
+
+from model.base import SelfForcingModel
+
+
+class DMD(SelfForcingModel):
+    def __init__(self, args, device):
+        """
+        Initialize the DMD (Distribution Matching Distillation) module.
+        This class is self-contained and compute generator and fake score losses
+        in the forward pass.
+        """
+        super().__init__(args, device)
+        self.num_frame_per_block = getattr(args, "num_frame_per_block", 1)
+        self.same_step_across_blocks = getattr(args, "same_step_across_blocks", True)
+        self.num_training_frames = getattr(args, "num_training_frames", 21)
+
+        if self.num_frame_per_block > 1:
+            self.generator.model.num_frame_per_block = self.num_frame_per_block
+
+        self.independent_first_frame = getattr(args, "independent_first_frame", False)
+        if self.independent_first_frame:
+            self.generator.model.independent_first_frame = True
+        if args.gradient_checkpointing:
+            self.generator.enable_gradient_checkpointing()
+            self.fake_score.enable_gradient_checkpointing()
+
+        # this will be init later with fsdp-wrapped modules
+        self.inference_pipeline: SelfForcingTrainingPipeline = None
+
+        # Step 2: Initialize all dmd hyperparameters
+        self.num_train_timestep = args.num_train_timestep
+        self.min_step = int(0.02 * self.num_train_timestep)
+        self.max_step = int(0.98 * self.num_train_timestep)
+        if hasattr(args, "real_guidance_scale"):
+            self.real_guidance_scale = args.real_guidance_scale
+            self.fake_guidance_scale = args.fake_guidance_scale
+        else:
+            self.real_guidance_scale = args.guidance_scale
+            self.fake_guidance_scale = 0.0
+        self.timestep_shift = getattr(args, "timestep_shift", 1.0)
+        self.ts_schedule = getattr(args, "ts_schedule", True)
+        self.ts_schedule_max = getattr(args, "ts_schedule_max", False)
+        self.min_score_timestep = getattr(args, "min_score_timestep", 0)
+
+        if getattr(self.scheduler, "alphas_cumprod", None) is not None:
+            self.scheduler.alphas_cumprod = self.scheduler.alphas_cumprod.to(device)
+        else:
+            self.scheduler.alphas_cumprod = None
+
+    def _compute_kl_grad(
+        self, noisy_image_or_video: torch.Tensor,
+        estimated_clean_image_or_video: torch.Tensor,
+        timestep: torch.Tensor,
+        conditional_dict: dict, unconditional_dict: dict,
+        normalization: bool = True
+    ) -> Tuple[torch.Tensor, dict]:
+        """
+        Compute the KL grad (eq 7 in https://arxiv.org/abs/2311.18828).
+        Input:
+            - noisy_image_or_video: a tensor with shape [B, F, C, H, W] where the number of frame is 1 for images.
+            - estimated_clean_image_or_video: a tensor with shape [B, F, C, H, W] representing the estimated clean image or video.
+            - timestep: a tensor with shape [B, F] containing the randomly generated timestep.
+            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
+            - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
+            - normalization: a boolean indicating whether to normalize the gradient.
+        Output:
+            - kl_grad: a tensor representing the KL grad.
+            - kl_log_dict: a dictionary containing the intermediate tensors for logging.
+        """
+        # Step 1: Compute the fake score
+        _, pred_fake_image_cond = self.fake_score(
+            noisy_image_or_video=noisy_image_or_video,
+            conditional_dict=conditional_dict,
+            timestep=timestep
+        )
+
+        if self.fake_guidance_scale != 0.0:
+            _, pred_fake_image_uncond = self.fake_score(
+                noisy_image_or_video=noisy_image_or_video,
+                conditional_dict=unconditional_dict,
+                timestep=timestep
+            )
+            pred_fake_image = pred_fake_image_cond + (
+                pred_fake_image_cond - pred_fake_image_uncond
+            ) * self.fake_guidance_scale
+        else:
+            pred_fake_image = pred_fake_image_cond
+
+        # Step 2: Compute the real score
+        # We compute the conditional and unconditional prediction
+        # and add them together to achieve cfg (https://arxiv.org/abs/2207.12598)
+        _, pred_real_image_cond = self.real_score(
+            noisy_image_or_video=noisy_image_or_video,
+            conditional_dict=conditional_dict,
+            timestep=timestep
+        )
+
+        _, pred_real_image_uncond = self.real_score(
+            noisy_image_or_video=noisy_image_or_video,
+            conditional_dict=unconditional_dict,
+            timestep=timestep
+        )
+
+        pred_real_image = pred_real_image_cond + (
+            pred_real_image_cond - pred_real_image_uncond
+        ) * self.real_guidance_scale
+
+        # Step 3: Compute the DMD gradient (DMD paper eq. 7).
+        grad = (pred_fake_image - pred_real_image)
+
+        # TODO: Change the normalizer for causal teacher
+        if normalization:
+            # Step 4: Gradient normalization (DMD paper eq. 8).
+            p_real = (estimated_clean_image_or_video - pred_real_image)
+            normalizer = torch.abs(p_real).mean(dim=[1, 2, 3, 4], keepdim=True)
+            grad = grad / normalizer
+        grad = torch.nan_to_num(grad)
+
+        return grad, {
+            "dmdtrain_gradient_norm": torch.mean(torch.abs(grad)).detach(),
+            "timestep": timestep.detach()
+        }
+
+    def compute_distribution_matching_loss(
+        self,
+        image_or_video: torch.Tensor,
+        conditional_dict: dict,
+        unconditional_dict: dict,
+        gradient_mask: Optional[torch.Tensor] = None,
+        denoised_timestep_from: int = 0,
+        denoised_timestep_to: int = 0
+    ) -> Tuple[torch.Tensor, dict]:
+        """
+        Compute the DMD loss (eq 7 in https://arxiv.org/abs/2311.18828).
+        Input:
+            - image_or_video: a tensor with shape [B, F, C, H, W] where the number of frame is 1 for images.
+            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
+            - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
+            - gradient_mask: a boolean tensor with the same shape as image_or_video indicating which pixels to compute loss .
+        Output:
+            - dmd_loss: a scalar tensor representing the DMD loss.
+            - dmd_log_dict: a dictionary containing the intermediate tensors for logging.
+        """
+        original_latent = image_or_video
+
+        batch_size, num_frame = image_or_video.shape[:2]
+
+        with torch.no_grad():
+            # Step 1: Randomly sample timestep based on the given schedule and corresponding noise
+            min_timestep = denoised_timestep_to if self.ts_schedule and denoised_timestep_to is not None else self.min_score_timestep
+            max_timestep = denoised_timestep_from if self.ts_schedule_max and denoised_timestep_from is not None else self.num_train_timestep
+            timestep = self._get_timestep(
+                min_timestep,
+                max_timestep,
+                batch_size,
+                num_frame,
+                self.num_frame_per_block,
+                uniform_timestep=True
+            )
+
+            # TODO:should we change it to `timestep = self.scheduler.timesteps[timestep]`?
+            if self.timestep_shift > 1:
+                timestep = self.timestep_shift * \
+                    (timestep / 1000) / \
+                    (1 + (self.timestep_shift - 1) * (timestep / 1000)) * 1000
+            timestep = timestep.clamp(self.min_step, self.max_step)
+
+            noise = torch.randn_like(image_or_video)
+            noisy_latent = self.scheduler.add_noise(
+                image_or_video.flatten(0, 1),
+                noise.flatten(0, 1),
+                timestep.flatten(0, 1)
+            ).detach().unflatten(0, (batch_size, num_frame))
+
+            # Step 2: Compute the KL grad
+            grad, dmd_log_dict = self._compute_kl_grad(
+                noisy_image_or_video=noisy_latent,
+                estimated_clean_image_or_video=original_latent,
+                timestep=timestep,
+                conditional_dict=conditional_dict,
+                unconditional_dict=unconditional_dict
+            )
+
+        if gradient_mask is not None:
+            dmd_loss = 0.5 * F.mse_loss(original_latent.double(
+            )[gradient_mask], (original_latent.double() - grad.double()).detach()[gradient_mask], reduction="mean")
+        else:
+            dmd_loss = 0.5 * F.mse_loss(original_latent.double(
+            ), (original_latent.double() - grad.double()).detach(), reduction="mean")
+        return dmd_loss, dmd_log_dict
+
+    def generator_loss(
+        self,
+        image_or_video_shape,
+        conditional_dict: dict,
+        unconditional_dict: dict,
+        clean_latent: torch.Tensor,
+        initial_latent: torch.Tensor = None
+    ) -> Tuple[torch.Tensor, dict]:
+        """
+        Generate image/videos from noise and compute the DMD loss.
+        The noisy input to the generator is backward simulated.
+        This removes the need of any datasets during distillation.
+        See Sec 4.5 of the DMD2 paper (https://arxiv.org/abs/2405.14867) for details.
+        Input:
+            - image_or_video_shape: a list containing the shape of the image or video [B, F, C, H, W].
+            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
+            - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
+            - clean_latent: a tensor containing the clean latents [B, F, C, H, W]. Need to be passed when no backward simulation is used.
+        Output:
+            - loss: a scalar tensor representing the generator loss.
+            - generator_log_dict: a dictionary containing the intermediate tensors for logging.
+        """
+        # Step 1: Unroll generator to obtain fake videos
+        pred_image, gradient_mask, denoised_timestep_from, denoised_timestep_to = self._run_generator(
+            image_or_video_shape=image_or_video_shape,
+            conditional_dict=conditional_dict,
+            initial_latent=initial_latent
+        )
+
+        # Step 2: Compute the DMD loss
+        dmd_loss, dmd_log_dict = self.compute_distribution_matching_loss(
+            image_or_video=pred_image,
+            conditional_dict=conditional_dict,
+            unconditional_dict=unconditional_dict,
+            gradient_mask=gradient_mask,
+            denoised_timestep_from=denoised_timestep_from,
+            denoised_timestep_to=denoised_timestep_to
+        )
+
+        return dmd_loss, dmd_log_dict
+
+    def critic_loss(
+        self,
+        image_or_video_shape,
+        conditional_dict: dict,
+        unconditional_dict: dict,
+        clean_latent: torch.Tensor,
+        initial_latent: torch.Tensor = None
+    ) -> Tuple[torch.Tensor, dict]:
+        """
+        Generate image/videos from noise and train the critic with generated samples.
+        The noisy input to the generator is backward simulated.
+        This removes the need of any datasets during distillation.
+        See Sec 4.5 of the DMD2 paper (https://arxiv.org/abs/2405.14867) for details.
+        Input:
+            - image_or_video_shape: a list containing the shape of the image or video [B, F, C, H, W].
+            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
+            - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
+            - clean_latent: a tensor containing the clean latents [B, F, C, H, W]. Need to be passed when no backward simulation is used.
+        Output:
+            - loss: a scalar tensor representing the generator loss.
+            - critic_log_dict: a dictionary containing the intermediate tensors for logging.
+        """
+
+        # Step 1: Run generator on backward simulated noisy input
+        with torch.no_grad():
+            generated_image, _, denoised_timestep_from, denoised_timestep_to = self._run_generator(
+                image_or_video_shape=image_or_video_shape,
+                conditional_dict=conditional_dict,
+                initial_latent=initial_latent
+            )
+
+        # Step 2: Compute the fake prediction
+        min_timestep = denoised_timestep_to if self.ts_schedule and denoised_timestep_to is not None else self.min_score_timestep
+        max_timestep = denoised_timestep_from if self.ts_schedule_max and denoised_timestep_from is not None else self.num_train_timestep
+        critic_timestep = self._get_timestep(
+            min_timestep,
+            max_timestep,
+            image_or_video_shape[0],
+            image_or_video_shape[1],
+            self.num_frame_per_block,
+            uniform_timestep=True
+        )
+
+        if self.timestep_shift > 1:
+            critic_timestep = self.timestep_shift * \
+                (critic_timestep / 1000) / (1 + (self.timestep_shift - 1) * (critic_timestep / 1000)) * 1000
+
+        critic_timestep = critic_timestep.clamp(self.min_step, self.max_step)
+
+        critic_noise = torch.randn_like(generated_image)
+        noisy_generated_image = self.scheduler.add_noise(
+            generated_image.flatten(0, 1),
+            critic_noise.flatten(0, 1),
+            critic_timestep.flatten(0, 1)
+        ).unflatten(0, image_or_video_shape[:2])
+
+        _, pred_fake_image = self.fake_score(
+            noisy_image_or_video=noisy_generated_image,
+            conditional_dict=conditional_dict,
+            timestep=critic_timestep
+        )
+
+        # Step 3: Compute the denoising loss for the fake critic
+        if self.args.denoising_loss_type == "flow":
+            from utils.wan_wrapper import WanDiffusionWrapper
+            flow_pred = WanDiffusionWrapper._convert_x0_to_flow_pred(
+                scheduler=self.scheduler,
+                x0_pred=pred_fake_image.flatten(0, 1),
+                xt=noisy_generated_image.flatten(0, 1),
+                timestep=critic_timestep.flatten(0, 1)
+            )
+            pred_fake_noise = None
+        else:
+            flow_pred = None
+            pred_fake_noise = self.scheduler.convert_x0_to_noise(
+                x0=pred_fake_image.flatten(0, 1),
+                xt=noisy_generated_image.flatten(0, 1),
+                timestep=critic_timestep.flatten(0, 1)
+            ).unflatten(0, image_or_video_shape[:2])
+
+        denoising_loss = self.denoising_loss_func(
+            x=generated_image.flatten(0, 1),
+            x_pred=pred_fake_image.flatten(0, 1),
+            noise=critic_noise.flatten(0, 1),
+            noise_pred=pred_fake_noise,
+            alphas_cumprod=self.scheduler.alphas_cumprod,
+            timestep=critic_timestep.flatten(0, 1),
+            flow_pred=flow_pred
+        )
+
+        # Step 5: Debugging Log
+        critic_log_dict = {
+            "critic_timestep": critic_timestep.detach()
+        }
+
+        return denoising_loss, critic_log_dict

model/gan.py

@@ -0,0 +1,295 @@
+import copy
+from pipeline import SelfForcingTrainingPipeline
+import torch.nn.functional as F
+from typing import Tuple
+import torch
+
+from model.base import SelfForcingModel
+
+
+class GAN(SelfForcingModel):
+    def __init__(self, args, device):
+        """
+        Initialize the GAN module.
+        This class is self-contained and compute generator and fake score losses
+        in the forward pass.
+        """
+        super().__init__(args, device)
+        self.num_frame_per_block = getattr(args, "num_frame_per_block", 1)
+        self.same_step_across_blocks = getattr(args, "same_step_across_blocks", True)
+        self.concat_time_embeddings = getattr(args, "concat_time_embeddings", False)
+        self.num_class = args.num_class
+        self.relativistic_discriminator = getattr(args, "relativistic_discriminator", False)
+
+        if self.num_frame_per_block > 1:
+            self.generator.model.num_frame_per_block = self.num_frame_per_block
+
+        self.fake_score.adding_cls_branch(
+            atten_dim=1536, num_class=args.num_class, time_embed_dim=1536 if self.concat_time_embeddings else 0)
+        self.fake_score.model.requires_grad_(True)
+
+        self.independent_first_frame = getattr(args, "independent_first_frame", False)
+        if self.independent_first_frame:
+            self.generator.model.independent_first_frame = True
+        if args.gradient_checkpointing:
+            self.generator.enable_gradient_checkpointing()
+            self.fake_score.enable_gradient_checkpointing()
+
+        # this will be init later with fsdp-wrapped modules
+        self.inference_pipeline: SelfForcingTrainingPipeline = None
+
+        # Step 2: Initialize all dmd hyperparameters
+        self.num_train_timestep = args.num_train_timestep
+        self.min_step = int(0.02 * self.num_train_timestep)
+        self.max_step = int(0.98 * self.num_train_timestep)
+        if hasattr(args, "real_guidance_scale"):
+            self.real_guidance_scale = args.real_guidance_scale
+            self.fake_guidance_scale = args.fake_guidance_scale
+        else:
+            self.real_guidance_scale = args.guidance_scale
+            self.fake_guidance_scale = 0.0
+        self.timestep_shift = getattr(args, "timestep_shift", 1.0)
+        self.critic_timestep_shift = getattr(args, "critic_timestep_shift", self.timestep_shift)
+        self.ts_schedule = getattr(args, "ts_schedule", True)
+        self.ts_schedule_max = getattr(args, "ts_schedule_max", False)
+        self.min_score_timestep = getattr(args, "min_score_timestep", 0)
+
+        self.gan_g_weight = getattr(args, "gan_g_weight", 1e-2)
+        self.gan_d_weight = getattr(args, "gan_d_weight", 1e-2)
+        self.r1_weight = getattr(args, "r1_weight", 0.0)
+        self.r2_weight = getattr(args, "r2_weight", 0.0)
+        self.r1_sigma = getattr(args, "r1_sigma", 0.01)
+        self.r2_sigma = getattr(args, "r2_sigma", 0.01)
+
+        if getattr(self.scheduler, "alphas_cumprod", None) is not None:
+            self.scheduler.alphas_cumprod = self.scheduler.alphas_cumprod.to(device)
+        else:
+            self.scheduler.alphas_cumprod = None
+
+    def _run_cls_pred_branch(self,
+                             noisy_image_or_video: torch.Tensor,
+                             conditional_dict: dict,
+                             timestep: torch.Tensor) -> torch.Tensor:
+        """
+            Run the classifier prediction branch on the generated image or video.
+            Input:
+                - image_or_video: a tensor with shape [B, F, C, H, W].
+            Output:
+                - cls_pred: a tensor with shape [B, 1, 1, 1, 1] representing the feature map for classification.
+        """
+        _, _, noisy_logit = self.fake_score(
+            noisy_image_or_video=noisy_image_or_video,
+            conditional_dict=conditional_dict,
+            timestep=timestep,
+            classify_mode=True,
+            concat_time_embeddings=self.concat_time_embeddings
+        )
+
+        return noisy_logit
+
+    def generator_loss(
+        self,
+        image_or_video_shape,
+        conditional_dict: dict,
+        unconditional_dict: dict,
+        clean_latent: torch.Tensor,
+        initial_latent: torch.Tensor = None
+    ) -> Tuple[torch.Tensor, dict]:
+        """
+        Generate image/videos from noise and compute the DMD loss.
+        The noisy input to the generator is backward simulated.
+        This removes the need of any datasets during distillation.
+        See Sec 4.5 of the DMD2 paper (https://arxiv.org/abs/2405.14867) for details.
+        Input:
+            - image_or_video_shape: a list containing the shape of the image or video [B, F, C, H, W].
+            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
+            - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
+            - clean_latent: a tensor containing the clean latents [B, F, C, H, W]. Need to be passed when no backward simulation is used.
+        Output:
+            - loss: a scalar tensor representing the generator loss.
+            - generator_log_dict: a dictionary containing the intermediate tensors for logging.
+        """
+        # Step 1: Unroll generator to obtain fake videos
+        pred_image, gradient_mask, denoised_timestep_from, denoised_timestep_to = self._run_generator(
+            image_or_video_shape=image_or_video_shape,
+            conditional_dict=conditional_dict,
+            initial_latent=initial_latent
+        )
+
+        # Step 2: Get timestep and add noise to generated/real latents
+        min_timestep = denoised_timestep_to if self.ts_schedule and denoised_timestep_to is not None else self.min_score_timestep
+        max_timestep = denoised_timestep_from if self.ts_schedule_max and denoised_timestep_from is not None else self.num_train_timestep
+        critic_timestep = self._get_timestep(
+            min_timestep,
+            max_timestep,
+            image_or_video_shape[0],
+            image_or_video_shape[1],
+            self.num_frame_per_block,
+            uniform_timestep=True
+        )
+
+        if self.critic_timestep_shift > 1:
+            critic_timestep = self.critic_timestep_shift * \
+                (critic_timestep / 1000) / (1 + (self.critic_timestep_shift - 1) * (critic_timestep / 1000)) * 1000
+
+        critic_timestep = critic_timestep.clamp(self.min_step, self.max_step)
+
+        critic_noise = torch.randn_like(pred_image)
+        noisy_fake_latent = self.scheduler.add_noise(
+            pred_image.flatten(0, 1),
+            critic_noise.flatten(0, 1),
+            critic_timestep.flatten(0, 1)
+        ).unflatten(0, image_or_video_shape[:2])
+
+        # Step 4: Compute the real GAN discriminator loss
+        real_image_or_video = clean_latent.clone()
+        critic_noise = torch.randn_like(real_image_or_video)
+        noisy_real_latent = self.scheduler.add_noise(
+            real_image_or_video.flatten(0, 1),
+            critic_noise.flatten(0, 1),
+            critic_timestep.flatten(0, 1)
+        ).unflatten(0, image_or_video_shape[:2])
+
+        conditional_dict["prompt_embeds"] = torch.concatenate(
+            (conditional_dict["prompt_embeds"], conditional_dict["prompt_embeds"]), dim=0)
+        critic_timestep = torch.concatenate((critic_timestep, critic_timestep), dim=0)
+        noisy_latent = torch.concatenate((noisy_fake_latent, noisy_real_latent), dim=0)
+        _, _, noisy_logit = self.fake_score(
+            noisy_image_or_video=noisy_latent,
+            conditional_dict=conditional_dict,
+            timestep=critic_timestep,
+            classify_mode=True,
+            concat_time_embeddings=self.concat_time_embeddings
+        )
+        noisy_fake_logit, noisy_real_logit = noisy_logit.chunk(2, dim=0)
+
+        if not self.relativistic_discriminator:
+            gan_G_loss = F.softplus(-noisy_fake_logit.float()).mean() * self.gan_g_weight
+        else:
+            relative_fake_logit = noisy_fake_logit - noisy_real_logit
+            gan_G_loss = F.softplus(-relative_fake_logit.float()).mean() * self.gan_g_weight
+
+        return gan_G_loss
+
+    def critic_loss(
+        self,
+        image_or_video_shape,
+        conditional_dict: dict,
+        unconditional_dict: dict,
+        clean_latent: torch.Tensor,
+        real_image_or_video: torch.Tensor,
+        initial_latent: torch.Tensor = None
+    ) -> Tuple[torch.Tensor, dict]:
+        """
+        Generate image/videos from noise and train the critic with generated samples.
+        The noisy input to the generator is backward simulated.
+        This removes the need of any datasets during distillation.
+        See Sec 4.5 of the DMD2 paper (https://arxiv.org/abs/2405.14867) for details.
+        Input:
+            - image_or_video_shape: a list containing the shape of the image or video [B, F, C, H, W].
+            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
+            - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
+            - clean_latent: a tensor containing the clean latents [B, F, C, H, W]. Need to be passed when no backward simulation is used.
+        Output:
+            - loss: a scalar tensor representing the generator loss.
+            - critic_log_dict: a dictionary containing the intermediate tensors for logging.
+        """
+
+        # Step 1: Run generator on backward simulated noisy input
+        with torch.no_grad():
+            generated_image, _, denoised_timestep_from, denoised_timestep_to, num_sim_steps = self._run_generator(
+                image_or_video_shape=image_or_video_shape,
+                conditional_dict=conditional_dict,
+                initial_latent=initial_latent
+            )
+
+        # Step 2: Get timestep and add noise to generated/real latents
+        min_timestep = denoised_timestep_to if self.ts_schedule and denoised_timestep_to is not None else self.min_score_timestep
+        max_timestep = denoised_timestep_from if self.ts_schedule_max and denoised_timestep_from is not None else self.num_train_timestep
+        critic_timestep = self._get_timestep(
+            min_timestep,
+            max_timestep,
+            image_or_video_shape[0],
+            image_or_video_shape[1],
+            self.num_frame_per_block,
+            uniform_timestep=True
+        )
+
+        if self.critic_timestep_shift > 1:
+            critic_timestep = self.critic_timestep_shift * \
+                (critic_timestep / 1000) / (1 + (self.critic_timestep_shift - 1) * (critic_timestep / 1000)) * 1000
+
+        critic_timestep = critic_timestep.clamp(self.min_step, self.max_step)
+
+        critic_noise = torch.randn_like(generated_image)
+        noisy_fake_latent = self.scheduler.add_noise(
+            generated_image.flatten(0, 1),
+            critic_noise.flatten(0, 1),
+            critic_timestep.flatten(0, 1)
+        ).unflatten(0, image_or_video_shape[:2])
+
+        # Step 4: Compute the real GAN discriminator loss
+        noisy_real_latent = self.scheduler.add_noise(
+            real_image_or_video.flatten(0, 1),
+            critic_noise.flatten(0, 1),
+            critic_timestep.flatten(0, 1)
+        ).unflatten(0, image_or_video_shape[:2])
+
+        conditional_dict_cloned = copy.deepcopy(conditional_dict)
+        conditional_dict_cloned["prompt_embeds"] = torch.concatenate(
+            (conditional_dict_cloned["prompt_embeds"], conditional_dict_cloned["prompt_embeds"]), dim=0)
+        _, _, noisy_logit = self.fake_score(
+            noisy_image_or_video=torch.concatenate((noisy_fake_latent, noisy_real_latent), dim=0),
+            conditional_dict=conditional_dict_cloned,
+            timestep=torch.concatenate((critic_timestep, critic_timestep), dim=0),
+            classify_mode=True,
+            concat_time_embeddings=self.concat_time_embeddings
+        )
+        noisy_fake_logit, noisy_real_logit = noisy_logit.chunk(2, dim=0)
+
+        if not self.relativistic_discriminator:
+            gan_D_loss = F.softplus(-noisy_real_logit.float()).mean() + F.softplus(noisy_fake_logit.float()).mean()
+        else:
+            relative_real_logit = noisy_real_logit - noisy_fake_logit
+            gan_D_loss = F.softplus(-relative_real_logit.float()).mean()
+        gan_D_loss = gan_D_loss * self.gan_d_weight
+
+        # R1 regularization
+        if self.r1_weight > 0.:
+            noisy_real_latent_perturbed = noisy_real_latent.clone()
+            epison_real = self.r1_sigma * torch.randn_like(noisy_real_latent_perturbed)
+            noisy_real_latent_perturbed = noisy_real_latent_perturbed + epison_real
+            noisy_real_logit_perturbed = self._run_cls_pred_branch(
+                noisy_image_or_video=noisy_real_latent_perturbed,
+                conditional_dict=conditional_dict,
+                timestep=critic_timestep
+            )
+
+            r1_grad = (noisy_real_logit_perturbed - noisy_real_logit) / self.r1_sigma
+            r1_loss = self.r1_weight * torch.mean((r1_grad)**2)
+        else:
+            r1_loss = torch.zeros_like(gan_D_loss)
+
+        # R2 regularization
+        if self.r2_weight > 0.:
+            noisy_fake_latent_perturbed = noisy_fake_latent.clone()
+            epison_generated = self.r2_sigma * torch.randn_like(noisy_fake_latent_perturbed)
+            noisy_fake_latent_perturbed = noisy_fake_latent_perturbed + epison_generated
+            noisy_fake_logit_perturbed = self._run_cls_pred_branch(
+                noisy_image_or_video=noisy_fake_latent_perturbed,
+                conditional_dict=conditional_dict,
+                timestep=critic_timestep
+            )
+
+            r2_grad = (noisy_fake_logit_perturbed - noisy_fake_logit) / self.r2_sigma
+            r2_loss = self.r2_weight * torch.mean((r2_grad)**2)
+        else:
+            r2_loss = torch.zeros_like(r2_loss)
+
+        critic_log_dict = {
+            "critic_timestep": critic_timestep.detach(),
+            'noisy_real_logit': noisy_real_logit.detach(),
+            'noisy_fake_logit': noisy_fake_logit.detach(),
+        }
+
+        return (gan_D_loss, r1_loss, r2_loss), critic_log_dict

model/ode_regression.py

@@ -0,0 +1,138 @@
+import torch.nn.functional as F
+from typing import Tuple
+import torch
+
+from model.base import BaseModel
+from utils.wan_wrapper import WanDiffusionWrapper, WanTextEncoder, WanVAEWrapper
+
+
+class ODERegression(BaseModel):
+    def __init__(self, args, device):
+        """
+        Initialize the ODERegression module.
+        This class is self-contained and compute generator losses
+        in the forward pass given precomputed ode solution pairs.
+        This class supports the ode regression loss for both causal and bidirectional models.
+        See Sec 4.3 of CausVid https://arxiv.org/abs/2412.07772 for details
+        """
+        super().__init__(args, device)
+
+        # Step 1: Initialize all models
+
+        self.generator = WanDiffusionWrapper(**getattr(args, "model_kwargs", {}), is_causal=True)
+        self.generator.model.requires_grad_(True)
+        if getattr(args, "generator_ckpt", False):
+            print(f"Loading pretrained generator from {args.generator_ckpt}")
+            state_dict = torch.load(args.generator_ckpt, map_location="cpu")[
+                'generator']
+            self.generator.load_state_dict(
+                state_dict, strict=True
+            )
+
+        self.num_frame_per_block = getattr(args, "num_frame_per_block", 1)
+
+        if self.num_frame_per_block > 1:
+            self.generator.model.num_frame_per_block = self.num_frame_per_block
+
+        self.independent_first_frame = getattr(args, "independent_first_frame", False)
+        if self.independent_first_frame:
+            self.generator.model.independent_first_frame = True
+        if args.gradient_checkpointing:
+            self.generator.enable_gradient_checkpointing()
+
+        # Step 2: Initialize all hyperparameters
+        self.timestep_shift = getattr(args, "timestep_shift", 1.0)
+
+    def _initialize_models(self, args):
+        self.generator = WanDiffusionWrapper(**getattr(args, "model_kwargs", {}), is_causal=True)
+        self.generator.model.requires_grad_(True)
+
+        self.text_encoder = WanTextEncoder()
+        self.text_encoder.requires_grad_(False)
+
+        self.vae = WanVAEWrapper()
+        self.vae.requires_grad_(False)
+
+    @torch.no_grad()
+    def _prepare_generator_input(self, ode_latent: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Given a tensor containing the whole ODE sampling trajectories,
+        randomly choose an intermediate timestep and return the latent as well as the corresponding timestep.
+        Input:
+            - ode_latent: a tensor containing the whole ODE sampling trajectories [batch_size, num_denoising_steps, num_frames, num_channels, height, width].
+        Output:
+            - noisy_input: a tensor containing the selected latent [batch_size, num_frames, num_channels, height, width].
+            - timestep: a tensor containing the corresponding timestep [batch_size].
+        """
+        batch_size, num_denoising_steps, num_frames, num_channels, height, width = ode_latent.shape
+
+        # Step 1: Randomly choose a timestep for each frame
+        index = self._get_timestep(
+            0,
+            len(self.denoising_step_list),
+            batch_size,
+            num_frames,
+            self.num_frame_per_block,
+            uniform_timestep=False
+        )
+        if self.args.i2v:
+            index[:, 0] = len(self.denoising_step_list) - 1
+
+        noisy_input = torch.gather(
+            ode_latent, dim=1,
+            index=index.reshape(batch_size, 1, num_frames, 1, 1, 1).expand(
+                -1, -1, -1, num_channels, height, width).to(self.device)
+        ).squeeze(1)
+
+        timestep = self.denoising_step_list[index].to(self.device)
+
+        # if self.extra_noise_step > 0:
+        #     random_timestep = torch.randint(0, self.extra_noise_step, [
+        #                                     batch_size, num_frames], device=self.device, dtype=torch.long)
+        #     perturbed_noisy_input = self.scheduler.add_noise(
+        #         noisy_input.flatten(0, 1),
+        #         torch.randn_like(noisy_input.flatten(0, 1)),
+        #         random_timestep.flatten(0, 1)
+        #     ).detach().unflatten(0, (batch_size, num_frames)).type_as(noisy_input)
+
+        #     noisy_input[timestep == 0] = perturbed_noisy_input[timestep == 0]
+
+        return noisy_input, timestep
+
+    def generator_loss(self, ode_latent: torch.Tensor, conditional_dict: dict) -> Tuple[torch.Tensor, dict]:
+        """
+        Generate image/videos from noisy latents and compute the ODE regression loss.
+        Input:
+            - ode_latent: a tensor containing the ODE latents [batch_size, num_denoising_steps, num_frames, num_channels, height, width].
+            They are ordered from most noisy to clean latents.
+            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
+        Output:
+            - loss: a scalar tensor representing the generator loss.
+            - log_dict: a dictionary containing additional information for loss timestep breakdown.
+        """
+        # Step 1: Run generator on noisy latents
+        target_latent = ode_latent[:, -1]
+
+        noisy_input, timestep = self._prepare_generator_input(
+            ode_latent=ode_latent)
+
+        _, pred_image_or_video = self.generator(
+            noisy_image_or_video=noisy_input,
+            conditional_dict=conditional_dict,
+            timestep=timestep
+        )
+
+        # Step 2: Compute the regression loss
+        mask = timestep != 0
+
+        loss = F.mse_loss(
+            pred_image_or_video[mask], target_latent[mask], reduction="mean")
+
+        log_dict = {
+            "unnormalized_loss": F.mse_loss(pred_image_or_video, target_latent, reduction='none').mean(dim=[1, 2, 3, 4]).detach(),
+            "timestep": timestep.float().mean(dim=1).detach(),
+            "input": noisy_input.detach(),
+            "output": pred_image_or_video.detach(),
+        }
+
+        return loss, log_dict

model/predictor_v4.py

@@ -0,0 +1,886 @@
+"""Wan Predictor-v4 for skipped Self-Forcing denoising steps.
+
+The Predictor is initialized from an already-loaded ``CausalWanModel``.  It
+keeps the Teacher patch/time/head modules frozen, trains two copied causal Wan
+blocks, and predicts a residual over the same-chunk anchor hidden state.
+
+Two history paths are supported:
+
+* online F-P-P-F inference can pass the generator's existing KV caches;
+* offline training can rebuild selected-layer history KV from clean-pass
+  self-attention prefeatures with :meth:`build_history_kv_cache`.
+
+The ordinary ``state_dict`` API is intentionally unchanged.  Use
+``trainable_state_dict``/``checkpoint_dict`` for compact Predictor checkpoints
+that omit frozen Teacher weights.
+"""
+
+from __future__ import annotations
+
+import copy
+import math
+from collections import OrderedDict
+from collections.abc import Mapping, Sequence
+from dataclasses import asdict, dataclass
+from typing import Any
+
+import torch
+from torch import nn
+
+from wan.modules.causal_model import (
+    CausalWanAttentionBlock,
+    CausalWanModel,
+    causal_rope_apply,
+)
+from wan.modules.model import sinusoidal_embedding_1d
+
+
+@dataclass(frozen=True)
+class WanPredictorV4Config:
+    """Serializable architecture metadata derived from the loaded Teacher."""
+
+    format_version: int
+    model_type: str
+    patch_size: tuple[int, int, int]
+    in_dim: int
+    dim: int
+    ffn_dim: int
+    freq_dim: int
+    out_dim: int
+    num_heads: int
+    num_layers: int
+    local_attn_size: int
+    sink_size: int
+    qk_norm: bool
+    cross_attn_norm: bool
+    eps: float
+    source_block_ids: tuple[int, int]
+    spatial_grid: tuple[int, int]
+
+    @property
+    def tokens_per_frame(self) -> int:
+        return math.prod(self.spatial_grid)
+
+    def to_dict(self) -> dict[str, Any]:
+        return asdict(self)
+
+
+class TripleFeatureFusion(nn.Module):
+    """Fuse target-latent, same-chunk anchor, and previous-chunk features."""
+
+    def __init__(self, dim: int, eps: float = 1e-6) -> None:
+        super().__init__()
+        self.current_norm = nn.LayerNorm(dim, eps=eps)
+        self.anchor_norm = nn.LayerNorm(dim, eps=eps)
+        self.previous_norm = nn.LayerNorm(dim, eps=eps)
+        self.mlp = nn.Sequential(
+            nn.Linear(3 * dim, 2 * dim),
+            nn.SiLU(),
+            nn.Linear(2 * dim, dim),
+        )
+
+    def forward(
+        self,
+        current: torch.Tensor,
+        anchor: torch.Tensor,
+        previous: torch.Tensor,
+    ) -> torch.Tensor:
+        if current.shape != anchor.shape or current.shape != previous.shape:
+            raise ValueError(
+                "TripleFeatureFusion requires identical [B, L, D] shapes, got "
+                f"current={tuple(current.shape)}, anchor={tuple(anchor.shape)}, "
+                f"previous={tuple(previous.shape)}"
+            )
+        return self.mlp(
+            torch.cat(
+                (
+                    self.current_norm(current),
+                    self.anchor_norm(anchor),
+                    self.previous_norm(previous),
+                ),
+                dim=-1,
+            )
+        )
+
+
+class _FrozenHistoryProjector(nn.Module):
+    """Frozen copy of one Teacher self-attention K/V projection path."""
+
+    def __init__(self, teacher_block: CausalWanAttentionBlock) -> None:
+        super().__init__()
+        self.k = copy.deepcopy(teacher_block.self_attn.k)
+        self.v = copy.deepcopy(teacher_block.self_attn.v)
+        self.norm_k = copy.deepcopy(teacher_block.self_attn.norm_k)
+        self.requires_grad_(False)
+
+    def forward(
+        self,
+        self_attn_input: torch.Tensor,
+        *,
+        num_heads: int,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        batch, tokens, dim = self_attn_input.shape
+        if dim % num_heads:
+            raise ValueError(f"Hidden dim {dim} is not divisible by {num_heads} heads")
+        head_dim = dim // num_heads
+        key = self.norm_k(self.k(self_attn_input)).view(
+            batch, tokens, num_heads, head_dim
+        )
+        value = self.v(self_attn_input).view(
+            batch, tokens, num_heads, head_dim
+        )
+        return key, value
+
+
+class SelfForcingPredictorV4(nn.Module):
+    """Two-block Wan Predictor used for the middle denoising steps of F-P-P-F."""
+
+    requires_previous_chunk_hidden = True
+    uses_history_kv = True
+    uses_clean_prefeature = True
+    checkpoint_format_version = 1
+
+    def __init__(
+        self,
+        teacher_model: CausalWanModel,
+        *,
+        source_block_ids: tuple[int, int] = (1, 28),
+        spatial_grid: tuple[int, int] = (30, 52),
+    ) -> None:
+        super().__init__()
+        teacher_model = self._unwrap_teacher(teacher_model)
+        if teacher_model.model_type != "t2v":
+            raise NotImplementedError("SelfForcingPredictorV4 currently supports Wan T2V")
+        if len(source_block_ids) != 2 or len(set(source_block_ids)) != 2:
+            raise ValueError("Predictor-v4 requires exactly two distinct source blocks")
+        if any(index < 0 or index >= len(teacher_model.blocks) for index in source_block_ids):
+            raise ValueError(
+                f"Invalid source blocks {source_block_ids} for "
+                f"{len(teacher_model.blocks)} Teacher blocks"
+            )
+        if len(spatial_grid) != 2 or any(int(size) <= 0 for size in spatial_grid):
+            raise ValueError(f"Invalid Predictor token spatial grid: {spatial_grid}")
+
+        first_self_attn = teacher_model.blocks[0].self_attn
+        self.predictor_config = WanPredictorV4Config(
+            format_version=self.checkpoint_format_version,
+            model_type=str(teacher_model.model_type),
+            patch_size=tuple(int(item) for item in teacher_model.patch_size),
+            in_dim=int(teacher_model.in_dim),
+            dim=int(teacher_model.dim),
+            ffn_dim=int(teacher_model.ffn_dim),
+            freq_dim=int(teacher_model.freq_dim),
+            out_dim=int(teacher_model.out_dim),
+            num_heads=int(teacher_model.num_heads),
+            num_layers=len(teacher_model.blocks),
+            local_attn_size=int(teacher_model.local_attn_size),
+            sink_size=int(first_self_attn.sink_size),
+            qk_norm=bool(teacher_model.qk_norm),
+            cross_attn_norm=bool(teacher_model.cross_attn_norm),
+            eps=float(teacher_model.eps),
+            source_block_ids=tuple(int(item) for item in source_block_ids),
+            spatial_grid=tuple(int(item) for item in spatial_grid),
+        )
+        cfg = self.predictor_config
+
+        # Frozen modules are copied rather than referenced so that calling
+        # Predictor.train()/to() cannot alter the loaded generator.
+        self.patch_embedding = copy.deepcopy(teacher_model.patch_embedding)
+        self.time_embedding = copy.deepcopy(teacher_model.time_embedding)
+        self.time_projection = copy.deepcopy(teacher_model.time_projection)
+        self.head = copy.deepcopy(teacher_model.head)
+        self.predictor_blocks = nn.ModuleList(
+            [copy.deepcopy(teacher_model.blocks[index]) for index in source_block_ids]
+        )
+        self.history_projectors = nn.ModuleDict(
+            {
+                str(index): _FrozenHistoryProjector(teacher_model.blocks[index])
+                for index in source_block_ids
+            }
+        )
+
+        self.feature_fusion = TripleFeatureFusion(cfg.dim, cfg.eps)
+        self.residual_out = nn.Linear(cfg.dim, cfg.dim)
+
+        self._freeze_teacher_modules()
+        self.predictor_blocks.requires_grad_(True)
+        # Text K/V come from the Teacher cross-attention cache.  Predictor only
+        # executes the query/output side, so keep unused cache-building weights
+        # frozen and out of the optimizer/checkpoint.
+        for block in self.predictor_blocks:
+            block.cross_attn.k.requires_grad_(False)
+            block.cross_attn.v.requires_grad_(False)
+            block.cross_attn.norm_k.requires_grad_(False)
+
+        reference = teacher_model.patch_embedding.weight
+        self.feature_fusion.to(device=reference.device, dtype=reference.dtype)
+        self.residual_out.to(device=reference.device, dtype=reference.dtype)
+        nn.init.zeros_(self.residual_out.weight)
+        nn.init.zeros_(self.residual_out.bias)
+
+        # Match CausalWanModel: RoPE frequencies are runtime state rather than a
+        # persistent buffer, so compact checkpoints contain no derived table.
+        self._freqs = teacher_model.freqs.detach().clone()
+
+    @staticmethod
+    def _unwrap_teacher(model: Any) -> CausalWanModel:
+        current = model
+        visited: set[int] = set()
+        while id(current) not in visited:
+            visited.add(id(current))
+            if isinstance(current, CausalWanModel):
+                return current
+            wrapped = getattr(current, "module", None)
+            if wrapped is not None:
+                current = wrapped
+                continue
+            nested = getattr(current, "model", None)
+            if nested is not None:
+                current = nested
+                continue
+            break
+        raise TypeError(
+            "teacher_model must be CausalWanModel (normally generator.model), "
+            f"got {type(model)!r}"
+        )
+
+    @classmethod
+    def from_teacher(
+        cls,
+        teacher_model: CausalWanModel,
+        *,
+        source_block_ids: tuple[int, int] = (1, 28),
+        spatial_grid: tuple[int, int] = (30, 52),
+    ) -> "SelfForcingPredictorV4":
+        """Initialize all copied/frozen/trainable weights from a loaded Teacher."""
+
+        return cls(
+            teacher_model,
+            source_block_ids=source_block_ids,
+            spatial_grid=spatial_grid,
+        )
+
+    @property
+    def config_dict(self) -> dict[str, Any]:
+        return self.predictor_config.to_dict()
+
+    @property
+    def source_block_ids(self) -> tuple[int, int]:
+        return self.predictor_config.source_block_ids
+
+    def _freeze_teacher_modules(self) -> None:
+        for module in (
+            self.patch_embedding,
+            self.time_embedding,
+            self.time_projection,
+            self.head,
+            self.history_projectors,
+        ):
+            module.requires_grad_(False)
+            module.eval()
+
+    @torch.no_grad()
+    def sync_frozen_from_teacher(
+        self,
+        teacher_model: CausalWanModel,
+    ) -> None:
+        """Refresh only the frozen Teacher-derived Predictor parameters.
+
+        Joint DMD changes the Full Generator after Predictor construction.  A
+        compact Predictor checkpoint is reconstructed from that updated Full
+        model at inference time, so the frozen training-time copies must track
+        it as well.  Predictor-owned trainable blocks/fusion are never
+        overwritten here.
+        """
+
+        teacher_model = self._unwrap_teacher(teacher_model)
+        for destination, source in (
+            (self.patch_embedding, teacher_model.patch_embedding),
+            (self.time_embedding, teacher_model.time_embedding),
+            (self.time_projection, teacher_model.time_projection),
+            (self.head, teacher_model.head),
+        ):
+            destination.load_state_dict(source.state_dict(), strict=True)
+
+        for position, source_id in enumerate(self.source_block_ids):
+            teacher_block = teacher_model.blocks[source_id]
+            while hasattr(teacher_block, "module"):
+                teacher_block = teacher_block.module
+            history = self.history_projectors[str(source_id)]
+            history.k.load_state_dict(
+                teacher_block.self_attn.k.state_dict(), strict=True
+            )
+            history.v.load_state_dict(
+                teacher_block.self_attn.v.state_dict(), strict=True
+            )
+            history.norm_k.load_state_dict(
+                teacher_block.self_attn.norm_k.state_dict(), strict=True
+            )
+
+            predictor_cross = self.predictor_blocks[position].cross_attn
+            teacher_cross = teacher_block.cross_attn
+            predictor_cross.k.load_state_dict(
+                teacher_cross.k.state_dict(), strict=True
+            )
+            predictor_cross.v.load_state_dict(
+                teacher_cross.v.state_dict(), strict=True
+            )
+            predictor_cross.norm_k.load_state_dict(
+                teacher_cross.norm_k.state_dict(), strict=True
+            )
+
+        self._freeze_teacher_modules()
+        for block in self.predictor_blocks:
+            block.cross_attn.k.requires_grad_(False)
+            block.cross_attn.v.requires_grad_(False)
+            block.cross_attn.norm_k.requires_grad_(False)
+
+    def train(self, mode: bool = True) -> "SelfForcingPredictorV4":
+        super().train(mode)
+        # Frozen layers have no stochastic operations today, but pinning their
+        # mode makes the intended boundary robust to future Wan changes.
+        for module in (
+            self.patch_embedding,
+            self.time_embedding,
+            self.time_projection,
+            self.head,
+            self.history_projectors,
+        ):
+            module.eval()
+        return self
+
+    def _runtime_freqs(self, device: torch.device) -> torch.Tensor:
+        if self._freqs.device != device:
+            self._freqs = self._freqs.to(device)
+        return self._freqs
+
+    @staticmethod
+    def _scalar_int(value: int | torch.Tensor, name: str) -> int:
+        if torch.is_tensor(value):
+            if value.numel() != 1:
+                raise ValueError(f"{name} must be scalar, got shape {tuple(value.shape)}")
+            value = value.detach().item()
+        result = int(value)
+        if result < 0:
+            raise ValueError(f"{name} must be non-negative, got {result}")
+        return result
+
+    @staticmethod
+    def _start_values(
+        value: int | torch.Tensor,
+        *,
+        batch: int,
+        name: str,
+    ) -> list[int]:
+        if torch.is_tensor(value):
+            values = [int(item) for item in value.detach().reshape(-1).cpu().tolist()]
+        else:
+            values = [int(value)]
+        if len(values) == 1:
+            values *= batch
+        if len(values) != batch:
+            raise ValueError(f"{name} has {len(values)} values for batch {batch}")
+        if any(item < 0 for item in values):
+            raise ValueError(f"{name} must contain non-negative frame indices")
+        return values
+
+    def _rope_history_key(
+        self,
+        key: torch.Tensor,
+        *,
+        start_frames: int | torch.Tensor,
+    ) -> torch.Tensor:
+        cfg = self.predictor_config
+        batch, tokens = key.shape[:2]
+        if tokens % cfg.tokens_per_frame:
+            raise ValueError(
+                f"History tokens {tokens} are not divisible by "
+                f"{cfg.tokens_per_frame} tokens/frame"
+            )
+        frames = tokens // cfg.tokens_per_frame
+        starts = self._start_values(start_frames, batch=batch, name="start_frames")
+        freqs = self._runtime_freqs(key.device)
+        grid = torch.tensor(
+            [[frames, *cfg.spatial_grid]],
+            dtype=torch.long,
+            device=key.device,
+        )
+        if len(set(starts)) == 1:
+            return causal_rope_apply(
+                key,
+                grid.expand(batch, -1),
+                freqs,
+                start_frame=starts[0],
+            )
+        return torch.cat(
+            [
+                causal_rope_apply(
+                    key[index : index + 1],
+                    grid,
+                    freqs,
+                    start_frame=start,
+                )
+                for index, start in enumerate(starts)
+            ],
+            dim=0,
+        )
+
+    def _project_history_part(
+        self,
+        block_id: int,
+        prefeature: torch.Tensor,
+        *,
+        start_frames: int | torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        cfg = self.predictor_config
+        if prefeature.ndim != 3 or prefeature.shape[-1] != cfg.dim:
+            raise ValueError(
+                f"Block {block_id} prefeature must be [B, S, {cfg.dim}], got "
+                f"{tuple(prefeature.shape)}"
+            )
+        projector = self.history_projectors[str(block_id)]
+        projector_device = projector.k.weight.device
+        if prefeature.device != projector_device:
+            raise ValueError(
+                f"Block {block_id} prefeature is on {prefeature.device}, "
+                f"projector is on {projector_device}"
+            )
+        prefeature = prefeature.to(dtype=projector.k.weight.dtype)
+        # Clean-history tensors are fixed offline Teacher data.  Avoid retaining
+        # a graph through several GiB of reconstructed cache.
+        with torch.no_grad():
+            key, value = projector(prefeature, num_heads=cfg.num_heads)
+            key = self._rope_history_key(key, start_frames=start_frames)
+        return key, value
+
+    def build_history_kv_cache(
+        self,
+        clean_prefeature_by_block: Mapping[
+            int | str, torch.Tensor | Sequence[torch.Tensor]
+        ],
+        *,
+        current_start: int | torch.Tensor,
+        current_tokens: int | torch.Tensor,
+        start_frames: int | torch.Tensor | Sequence[int | torch.Tensor] = 0,
+        cache_capacity: int | None = None,
+    ) -> dict[int, dict[str, torch.Tensor]]:
+        """Rebuild selected-layer clean-history caches for Predictor training.
+
+        ``clean_prefeature_by_block`` may contain one already-concatenated
+        ``[B, S, D]`` tensor per block, or a sequence of chunk tensors.  For a
+        sequence, ``start_frames`` can be the matching sequence ``0, 3, ...``.
+        ``current_start`` is the global token offset used by Wan inference and
+        ``current_tokens`` reserves the writable current-chunk cache region.
+        """
+
+        current_start_int = self._scalar_int(current_start, "current_start")
+        current_tokens_int = self._scalar_int(current_tokens, "current_tokens")
+        if current_tokens_int == 0:
+            raise ValueError("current_tokens must be positive")
+
+        cfg = self.predictor_config
+        result: dict[int, dict[str, torch.Tensor]] = {}
+        expected_batch: int | None = None
+        expected_history_tokens: int | None = None
+        for block_id in self.source_block_ids:
+            value = clean_prefeature_by_block.get(block_id)
+            if value is None:
+                value = clean_prefeature_by_block.get(str(block_id))
+            if value is None:
+                raise ValueError(f"Missing clean prefeature for block {block_id}")
+
+            if torch.is_tensor(value):
+                if isinstance(start_frames, Sequence) and not torch.is_tensor(start_frames):
+                    if len(start_frames) != 1:
+                        raise ValueError(
+                            "Already-concatenated prefeatures require one start_frames value"
+                        )
+                    part_start = start_frames[0]
+                else:
+                    part_start = start_frames
+                keys, values = self._project_history_part(
+                    block_id, value, start_frames=part_start
+                )
+            else:
+                parts = list(value)
+                if not parts:
+                    raise ValueError(f"Block {block_id} has no history prefeatures")
+                if isinstance(start_frames, Sequence) and not torch.is_tensor(start_frames):
+                    starts = list(start_frames)
+                    if len(starts) != len(parts):
+                        raise ValueError(
+                            f"start_frames has {len(starts)} entries for "
+                            f"{len(parts)} history chunks"
+                        )
+                else:
+                    starts = []
+                    next_start: int | torch.Tensor = start_frames
+                    for part in parts:
+                        starts.append(next_start)
+                        if torch.is_tensor(next_start) and next_start.numel() > 1:
+                            next_start = next_start + (
+                                part.shape[1] // self.predictor_config.tokens_per_frame
+                            )
+                        else:
+                            next_start = self._scalar_int(next_start, "start_frames") + (
+                                part.shape[1] // self.predictor_config.tokens_per_frame
+                            )
+                projected = [
+                    self._project_history_part(
+                        block_id, part, start_frames=part_start
+                    )
+                    for part, part_start in zip(parts, starts)
+                ]
+                keys = torch.cat([item[0] for item in projected], dim=1)
+                values = torch.cat([item[1] for item in projected], dim=1)
+
+            batch, history_tokens = keys.shape[:2]
+            if expected_batch is None:
+                expected_batch = batch
+                expected_history_tokens = history_tokens
+            elif batch != expected_batch or history_tokens != expected_history_tokens:
+                raise ValueError(
+                    "Selected blocks must have the same history shape, got "
+                    f"block {block_id}: batch={batch}, tokens={history_tokens}; "
+                    f"expected batch={expected_batch}, tokens={expected_history_tokens}"
+                )
+            if history_tokens > current_start_int:
+                raise ValueError(
+                    f"History has {history_tokens} tokens but current_start is "
+                    f"{current_start_int}"
+                )
+
+            required_capacity = history_tokens + current_tokens_int
+            capacity = required_capacity if cache_capacity is None else int(cache_capacity)
+            if capacity < required_capacity:
+                raise ValueError(
+                    f"cache_capacity {capacity} is smaller than required "
+                    f"{required_capacity}"
+                )
+            cache_k = keys.new_zeros(
+                batch, capacity, cfg.num_heads, cfg.dim // cfg.num_heads
+            )
+            cache_v = values.new_zeros(
+                batch, capacity, cfg.num_heads, cfg.dim // cfg.num_heads
+            )
+            cache_k[:, :history_tokens].copy_(keys)
+            cache_v[:, :history_tokens].copy_(values)
+            result[block_id] = {
+                "k": cache_k,
+                "v": cache_v,
+                "global_end_index": torch.tensor(
+                    [current_start_int], dtype=torch.long, device=keys.device
+                ),
+                "local_end_index": torch.tensor(
+                    [history_tokens], dtype=torch.long, device=keys.device
+                ),
+            }
+        return result
+
+    def _select_cache(
+        self,
+        caches: Mapping[Any, Any] | Sequence[Any],
+        *,
+        source_id: int,
+        source_position: int,
+        name: str,
+    ) -> Mapping[str, Any]:
+        if isinstance(caches, Mapping):
+            selected = caches.get(source_id)
+            if selected is None:
+                selected = caches.get(str(source_id))
+        else:
+            if len(caches) == self.predictor_config.num_layers:
+                selected = caches[source_id]
+            elif len(caches) == len(self.source_block_ids):
+                selected = caches[source_position]
+            else:
+                selected = None
+        if selected is None:
+            raise ValueError(f"{name} is missing source block {source_id}")
+        if not isinstance(selected, Mapping):
+            raise TypeError(f"{name}[{source_id}] must be a mapping")
+        return selected
+
+    def _selected_crossattn_cache(
+        self,
+        caches: Mapping[Any, Any] | Sequence[Any],
+        *,
+        source_id: int,
+        source_position: int,
+    ) -> dict[str, Any]:
+        selected = self._select_cache(
+            caches,
+            source_id=source_id,
+            source_position=source_position,
+            name="crossattn_cache",
+        )
+        missing = {"k", "v"}.difference(selected)
+        if missing:
+            raise ValueError(
+                f"crossattn_cache block {source_id} is missing {sorted(missing)}"
+            )
+        if selected["k"].shape != selected["v"].shape:
+            raise ValueError(f"crossattn_cache block {source_id} K/V shape mismatch")
+        # Offline files contain K/V but need not serialize the runtime flag.
+        # A shallow wrapper avoids mutating the generator-owned dictionary.
+        return {**selected, "is_init": True}
+
+    def _time_condition(
+        self,
+        target_timestep: torch.Tensor,
+        reference: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        cfg = self.predictor_config
+        time_embedding = self.time_embedding(
+            sinusoidal_embedding_1d(
+                cfg.freq_dim, target_timestep.flatten()
+            ).type_as(reference)
+        )
+        block_condition = self.time_projection(time_embedding).unflatten(
+            1, (6, cfg.dim)
+        ).unflatten(0, target_timestep.shape)
+        head_condition = time_embedding.unflatten(
+            0, target_timestep.shape
+        ).unsqueeze(2)
+        return block_condition, head_condition
+
+    def _unpatchify(
+        self,
+        tokens: torch.Tensor,
+        grid_sizes: torch.Tensor,
+    ) -> torch.Tensor:
+        cfg = self.predictor_config
+        outputs = []
+        for sample, grid in zip(tokens, grid_sizes.tolist()):
+            sample = sample[: math.prod(grid)].view(
+                *grid, *cfg.patch_size, cfg.out_dim
+            )
+            # [f, h, w, p, q, r, c] -> [f, p, c, h, q, w, r] so the
+            # following reshape merges each grid axis with its patch axis.
+            sample = torch.einsum("fhwpqrc->fpchqwr", sample)
+            outputs.append(
+                sample.reshape(
+                    grid[0] * cfg.patch_size[0],
+                    cfg.out_dim,
+                    grid[1] * cfg.patch_size[1],
+                    grid[2] * cfg.patch_size[2],
+                )
+            )
+        return torch.stack(outputs)
+
+    def forward(
+        self,
+        *,
+        target_latent: torch.Tensor,
+        target_timestep: torch.Tensor,
+        anchor_hidden: torch.Tensor,
+        previous_chunk_hidden: torch.Tensor,
+        kv_cache: Mapping[Any, Any] | Sequence[Any],
+        crossattn_cache: Mapping[Any, Any] | Sequence[Any],
+        current_start: int | torch.Tensor,
+    ) -> dict[str, torch.Tensor]:
+        """Predict one skipped denoising step.
+
+        Args:
+            target_latent: Noisy target chunk in ``[B, F, C, H, W]`` layout.
+            target_timestep: Per-frame timestep tensor ``[B, F]``.
+            anchor_hidden: Same-chunk preceding-step final hidden ``[B, L, D]``.
+            previous_chunk_hidden: Previous-chunk same-step hidden ``[B, L, D]``.
+            kv_cache: Full 30-layer list or selected-layer mapping/list.
+            crossattn_cache: Full list or selected cached text K/V.
+            current_start: Current chunk's global token offset.
+        """
+
+        cfg = self.predictor_config
+        if target_latent.ndim != 5:
+            raise ValueError(
+                f"target_latent must be [B, F, C, H, W], got {tuple(target_latent.shape)}"
+            )
+        batch, frames, channels, height, width = target_latent.shape
+        if channels != cfg.in_dim:
+            raise ValueError(f"target_latent channels {channels} != {cfg.in_dim}")
+        expected_timestep = (batch, frames // cfg.patch_size[0])
+        if tuple(target_timestep.shape) != expected_timestep:
+            raise ValueError(
+                f"target_timestep shape {tuple(target_timestep.shape)} != "
+                f"{expected_timestep}"
+            )
+        if target_timestep.device != target_latent.device:
+            raise ValueError("target_timestep and target_latent must share a device")
+        if target_latent.device != self.patch_embedding.weight.device:
+            raise ValueError(
+                f"target_latent is on {target_latent.device}, Predictor is on "
+                f"{self.patch_embedding.weight.device}"
+            )
+
+        latent_cf = target_latent.permute(0, 2, 1, 3, 4).to(
+            dtype=self.patch_embedding.weight.dtype
+        )
+        # Do not wrap frozen patch/time/head modules in no_grad: recursive
+        # rollout losses must still backpropagate to an earlier target latent.
+        current = self.patch_embedding(latent_cf)
+        grid_sizes = torch.tensor(
+            [current.shape[2:]] * batch,
+            dtype=torch.long,
+            device=current.device,
+        )
+        current = current.flatten(2).transpose(1, 2)
+        expected_hidden = (batch, current.shape[1], cfg.dim)
+        if tuple(anchor_hidden.shape) != expected_hidden:
+            raise ValueError(
+                f"anchor_hidden shape {tuple(anchor_hidden.shape)} != {expected_hidden}"
+            )
+        if tuple(previous_chunk_hidden.shape) != expected_hidden:
+            raise ValueError(
+                "previous_chunk_hidden shape "
+                f"{tuple(previous_chunk_hidden.shape)} != {expected_hidden}"
+            )
+        current = current.to(dtype=anchor_hidden.dtype)
+        hidden = self.feature_fusion(
+            current, anchor_hidden, previous_chunk_hidden
+        )
+        block_condition, head_condition = self._time_condition(
+            target_timestep, current
+        )
+        seq_lens = torch.full(
+            (batch,), current.shape[1], dtype=torch.long, device=current.device
+        )
+        current_start_int = self._scalar_int(current_start, "current_start")
+        freqs = self._runtime_freqs(current.device)
+
+        for position, (source_id, block) in enumerate(
+            zip(self.source_block_ids, self.predictor_blocks)
+        ):
+            selected_kv = self._select_cache(
+                kv_cache,
+                source_id=source_id,
+                source_position=position,
+                name="kv_cache",
+            )
+            selected_cross = self._selected_crossattn_cache(
+                crossattn_cache,
+                source_id=source_id,
+                source_position=position,
+            )
+            hidden = block(
+                hidden,
+                e=block_condition,
+                seq_lens=seq_lens,
+                grid_sizes=grid_sizes,
+                freqs=freqs,
+                context=None,
+                context_lens=None,
+                block_mask=None,
+                kv_cache=selected_kv,
+                crossattn_cache=selected_cross,
+                current_start=current_start_int,
+                cache_start=current_start_int,
+            )
+
+        delta_hidden = self.residual_out(hidden)
+        pred_hidden = anchor_hidden + delta_hidden
+        pred_tokens = self.head(pred_hidden, head_condition)
+        pred_flow = self._unpatchify(pred_tokens, grid_sizes)
+        return {
+            "pred_hidden": pred_hidden,
+            "pred_flow": pred_flow,
+            "delta_hidden": delta_hidden,
+        }
+
+    def trainable_parameter_count(self) -> int:
+        return sum(
+            parameter.numel()
+            for parameter in self.parameters()
+            if parameter.requires_grad
+        )
+
+    def trainable_parameter_breakdown(self) -> dict[str, int]:
+        modules = {
+            "feature_fusion": self.feature_fusion,
+            "predictor_blocks": self.predictor_blocks,
+            "residual_out": self.residual_out,
+        }
+        return {
+            name: sum(
+                parameter.numel()
+                for parameter in module.parameters()
+                if parameter.requires_grad
+            )
+            for name, module in modules.items()
+        }
+
+    def trainable_state_dict(
+        self,
+        *,
+        keep_vars: bool = False,
+    ) -> OrderedDict[str, torch.Tensor]:
+        """Return only optimizer-owned parameters, suitable for safetensors."""
+
+        trainable = {
+            name for name, parameter in self.named_parameters()
+            if parameter.requires_grad
+        }
+        state = super().state_dict(keep_vars=keep_vars)
+        return OrderedDict(
+            (name, value) for name, value in state.items() if name in trainable
+        )
+
+    def load_trainable_state_dict(
+        self,
+        state_dict: Mapping[str, torch.Tensor],
+        *,
+        strict: bool = True,
+    ) -> None:
+        """Load a compact state into a fresh Predictor initialized from Teacher."""
+
+        expected = set(self.trainable_state_dict())
+        received = set(state_dict)
+        if strict:
+            missing = sorted(expected.difference(received))
+            unexpected = sorted(received.difference(expected))
+            if missing or unexpected:
+                raise RuntimeError(
+                    "Predictor trainable checkpoint mismatch: "
+                    f"missing={missing}, unexpected={unexpected}"
+                )
+        filtered = {
+            name: tensor for name, tensor in state_dict.items() if name in expected
+        }
+        self.load_state_dict(filtered, strict=False)
+
+    def checkpoint_dict(self) -> dict[str, Any]:
+        """Build a compact torch-save payload with architecture metadata."""
+
+        return {
+            "format": "self_forcing_wan_predictor_v4",
+            "format_version": self.checkpoint_format_version,
+            "config": self.config_dict,
+            "trainable_state_dict": self.trainable_state_dict(),
+        }
+
+    def load_checkpoint_dict(
+        self,
+        checkpoint: Mapping[str, Any],
+        *,
+        strict: bool = True,
+    ) -> None:
+        if checkpoint.get("format") != "self_forcing_wan_predictor_v4":
+            raise ValueError(f"Unsupported Predictor checkpoint: {checkpoint.get('format')}")
+        saved_config = dict(checkpoint.get("config", {}))
+        if strict and saved_config != self.config_dict:
+            raise ValueError(
+                "Predictor checkpoint config does not match the Teacher/config "
+                f"used for reconstruction: saved={saved_config}, current={self.config_dict}"
+            )
+        state = checkpoint.get("trainable_state_dict")
+        if not isinstance(state, Mapping):
+            raise ValueError("Predictor checkpoint has no trainable_state_dict")
+        self.load_trainable_state_dict(state, strict=strict)
+
+
+__all__ = [
+    "SelfForcingPredictorV4",
+    "TripleFeatureFusion",
+    "WanPredictorV4Config",
+]

model/sid.py

@@ -0,0 +1,283 @@
+from pipeline import SelfForcingTrainingPipeline
+from typing import Optional, Tuple
+import torch
+
+from model.base import SelfForcingModel
+
+
+class SiD(SelfForcingModel):
+    def __init__(self, args, device):
+        """
+        Initialize the DMD (Distribution Matching Distillation) module.
+        This class is self-contained and compute generator and fake score losses
+        in the forward pass.
+        """
+        super().__init__(args, device)
+        self.num_frame_per_block = getattr(args, "num_frame_per_block", 1)
+
+        if self.num_frame_per_block > 1:
+            self.generator.model.num_frame_per_block = self.num_frame_per_block
+
+        if args.gradient_checkpointing:
+            self.generator.enable_gradient_checkpointing()
+            self.fake_score.enable_gradient_checkpointing()
+            self.real_score.enable_gradient_checkpointing()
+
+        # this will be init later with fsdp-wrapped modules
+        self.inference_pipeline: SelfForcingTrainingPipeline = None
+
+        # Step 2: Initialize all dmd hyperparameters
+        self.num_train_timestep = args.num_train_timestep
+        self.min_step = int(0.02 * self.num_train_timestep)
+        self.max_step = int(0.98 * self.num_train_timestep)
+        if hasattr(args, "real_guidance_scale"):
+            self.real_guidance_scale = args.real_guidance_scale
+        else:
+            self.real_guidance_scale = args.guidance_scale
+        self.timestep_shift = getattr(args, "timestep_shift", 1.0)
+        self.sid_alpha = getattr(args, "sid_alpha", 1.0)
+        self.ts_schedule = getattr(args, "ts_schedule", True)
+        self.ts_schedule_max = getattr(args, "ts_schedule_max", False)
+
+        if getattr(self.scheduler, "alphas_cumprod", None) is not None:
+            self.scheduler.alphas_cumprod = self.scheduler.alphas_cumprod.to(device)
+        else:
+            self.scheduler.alphas_cumprod = None
+
+    def compute_distribution_matching_loss(
+        self,
+        image_or_video: torch.Tensor,
+        conditional_dict: dict,
+        unconditional_dict: dict,
+        gradient_mask: Optional[torch.Tensor] = None,
+        denoised_timestep_from: int = 0,
+        denoised_timestep_to: int = 0
+    ) -> Tuple[torch.Tensor, dict]:
+        """
+        Compute the DMD loss (eq 7 in https://arxiv.org/abs/2311.18828).
+        Input:
+            - image_or_video: a tensor with shape [B, F, C, H, W] where the number of frame is 1 for images.
+            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
+            - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
+            - gradient_mask: a boolean tensor with the same shape as image_or_video indicating which pixels to compute loss .
+        Output:
+            - dmd_loss: a scalar tensor representing the DMD loss.
+            - dmd_log_dict: a dictionary containing the intermediate tensors for logging.
+        """
+        original_latent = image_or_video
+
+        batch_size, num_frame = image_or_video.shape[:2]
+
+        # Step 1: Randomly sample timestep based on the given schedule and corresponding noise
+        min_timestep = denoised_timestep_to if self.ts_schedule and denoised_timestep_to is not None else self.min_score_timestep
+        max_timestep = denoised_timestep_from if self.ts_schedule_max and denoised_timestep_from is not None else self.num_train_timestep
+        timestep = self._get_timestep(
+            min_timestep,
+            max_timestep,
+            batch_size,
+            num_frame,
+            self.num_frame_per_block,
+            uniform_timestep=True
+        )
+
+        if self.timestep_shift > 1:
+            timestep = self.timestep_shift * \
+                (timestep / 1000) / \
+                (1 + (self.timestep_shift - 1) * (timestep / 1000)) * 1000
+        timestep = timestep.clamp(self.min_step, self.max_step)
+
+        noise = torch.randn_like(image_or_video)
+        noisy_latent = self.scheduler.add_noise(
+            image_or_video.flatten(0, 1),
+            noise.flatten(0, 1),
+            timestep.flatten(0, 1)
+        ).unflatten(0, (batch_size, num_frame))
+
+        # Step 2: SiD (May be wrap it?)
+        noisy_image_or_video = noisy_latent
+        # Step 2.1: Compute the fake score
+        _, pred_fake_image = self.fake_score(
+            noisy_image_or_video=noisy_image_or_video,
+            conditional_dict=conditional_dict,
+            timestep=timestep
+        )
+        # Step 2.2: Compute the real score
+        # We compute the conditional and unconditional prediction
+        # and add them together to achieve cfg (https://arxiv.org/abs/2207.12598)
+        # NOTE: This step may cause OOM issue, which can be addressed by the CFG-free technique
+
+        _, pred_real_image_cond = self.real_score(
+            noisy_image_or_video=noisy_image_or_video,
+            conditional_dict=conditional_dict,
+            timestep=timestep
+        )
+
+        _, pred_real_image_uncond = self.real_score(
+            noisy_image_or_video=noisy_image_or_video,
+            conditional_dict=unconditional_dict,
+            timestep=timestep
+        )
+
+        pred_real_image = pred_real_image_cond + (
+            pred_real_image_cond - pred_real_image_uncond
+        ) * self.real_guidance_scale
+
+        # Step 2.3: SiD Loss
+        # TODO: Add alpha
+        # TODO: Double?
+        sid_loss = (pred_real_image.double() - pred_fake_image.double()) * ((pred_real_image.double() - original_latent.double()) - self.sid_alpha * (pred_real_image.double() - pred_fake_image.double()))
+
+        # Step 2.4: Loss normalizer
+        with torch.no_grad():
+            p_real = (original_latent - pred_real_image)
+            normalizer = torch.abs(p_real).mean(dim=[1, 2, 3, 4], keepdim=True)
+        sid_loss = sid_loss / normalizer
+
+        sid_loss = torch.nan_to_num(sid_loss)
+        num_frame = sid_loss.shape[1]
+        sid_loss = sid_loss.mean()
+
+        sid_log_dict = {
+            "dmdtrain_gradient_norm": torch.zeros_like(sid_loss),
+            "timestep": timestep.detach()
+        }
+
+        return sid_loss, sid_log_dict
+
+    def generator_loss(
+        self,
+        image_or_video_shape,
+        conditional_dict: dict,
+        unconditional_dict: dict,
+        clean_latent: torch.Tensor,
+        initial_latent: torch.Tensor = None
+    ) -> Tuple[torch.Tensor, dict]:
+        """
+        Generate image/videos from noise and compute the DMD loss.
+        The noisy input to the generator is backward simulated.
+        This removes the need of any datasets during distillation.
+        See Sec 4.5 of the DMD2 paper (https://arxiv.org/abs/2405.14867) for details.
+        Input:
+            - image_or_video_shape: a list containing the shape of the image or video [B, F, C, H, W].
+            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
+            - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
+            - clean_latent: a tensor containing the clean latents [B, F, C, H, W]. Need to be passed when no backward simulation is used.
+        Output:
+            - loss: a scalar tensor representing the generator loss.
+            - generator_log_dict: a dictionary containing the intermediate tensors for logging.
+        """
+        # Step 1: Unroll generator to obtain fake videos
+        pred_image, gradient_mask, denoised_timestep_from, denoised_timestep_to = self._run_generator(
+            image_or_video_shape=image_or_video_shape,
+            conditional_dict=conditional_dict,
+            initial_latent=initial_latent
+        )
+
+        # Step 2: Compute the DMD loss
+        dmd_loss, dmd_log_dict = self.compute_distribution_matching_loss(
+            image_or_video=pred_image,
+            conditional_dict=conditional_dict,
+            unconditional_dict=unconditional_dict,
+            gradient_mask=gradient_mask,
+            denoised_timestep_from=denoised_timestep_from,
+            denoised_timestep_to=denoised_timestep_to
+        )
+
+        return dmd_loss, dmd_log_dict
+
+    def critic_loss(
+        self,
+        image_or_video_shape,
+        conditional_dict: dict,
+        unconditional_dict: dict,
+        clean_latent: torch.Tensor,
+        initial_latent: torch.Tensor = None
+    ) -> Tuple[torch.Tensor, dict]:
+        """
+        Generate image/videos from noise and train the critic with generated samples.
+        The noisy input to the generator is backward simulated.
+        This removes the need of any datasets during distillation.
+        See Sec 4.5 of the DMD2 paper (https://arxiv.org/abs/2405.14867) for details.
+        Input:
+            - image_or_video_shape: a list containing the shape of the image or video [B, F, C, H, W].
+            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
+            - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
+            - clean_latent: a tensor containing the clean latents [B, F, C, H, W]. Need to be passed when no backward simulation is used.
+        Output:
+            - loss: a scalar tensor representing the generator loss.
+            - critic_log_dict: a dictionary containing the intermediate tensors for logging.
+        """
+
+        # Step 1: Run generator on backward simulated noisy input
+        with torch.no_grad():
+            generated_image, _, denoised_timestep_from, denoised_timestep_to = self._run_generator(
+                image_or_video_shape=image_or_video_shape,
+                conditional_dict=conditional_dict,
+                initial_latent=initial_latent
+            )
+
+        # Step 2: Compute the fake prediction
+        min_timestep = denoised_timestep_to if self.ts_schedule and denoised_timestep_to is not None else self.min_score_timestep
+        max_timestep = denoised_timestep_from if self.ts_schedule_max and denoised_timestep_from is not None else self.num_train_timestep
+        critic_timestep = self._get_timestep(
+            min_timestep,
+            max_timestep,
+            image_or_video_shape[0],
+            image_or_video_shape[1],
+            self.num_frame_per_block,
+            uniform_timestep=True
+        )
+
+        if self.timestep_shift > 1:
+            critic_timestep = self.timestep_shift * \
+                (critic_timestep / 1000) / (1 + (self.timestep_shift - 1) * (critic_timestep / 1000)) * 1000
+
+        critic_timestep = critic_timestep.clamp(self.min_step, self.max_step)
+
+        critic_noise = torch.randn_like(generated_image)
+        noisy_generated_image = self.scheduler.add_noise(
+            generated_image.flatten(0, 1),
+            critic_noise.flatten(0, 1),
+            critic_timestep.flatten(0, 1)
+        ).unflatten(0, image_or_video_shape[:2])
+
+        _, pred_fake_image = self.fake_score(
+            noisy_image_or_video=noisy_generated_image,
+            conditional_dict=conditional_dict,
+            timestep=critic_timestep
+        )
+
+        # Step 3: Compute the denoising loss for the fake critic
+        if self.args.denoising_loss_type == "flow":
+            from utils.wan_wrapper import WanDiffusionWrapper
+            flow_pred = WanDiffusionWrapper._convert_x0_to_flow_pred(
+                scheduler=self.scheduler,
+                x0_pred=pred_fake_image.flatten(0, 1),
+                xt=noisy_generated_image.flatten(0, 1),
+                timestep=critic_timestep.flatten(0, 1)
+            )
+            pred_fake_noise = None
+        else:
+            flow_pred = None
+            pred_fake_noise = self.scheduler.convert_x0_to_noise(
+                x0=pred_fake_image.flatten(0, 1),
+                xt=noisy_generated_image.flatten(0, 1),
+                timestep=critic_timestep.flatten(0, 1)
+            ).unflatten(0, image_or_video_shape[:2])
+
+        denoising_loss = self.denoising_loss_func(
+            x=generated_image.flatten(0, 1),
+            x_pred=pred_fake_image.flatten(0, 1),
+            noise=critic_noise.flatten(0, 1),
+            noise_pred=pred_fake_noise,
+            alphas_cumprod=self.scheduler.alphas_cumprod,
+            timestep=critic_timestep.flatten(0, 1),
+            flow_pred=flow_pred
+        )
+
+        # Step 5: Debugging Log
+        critic_log_dict = {
+            "critic_timestep": critic_timestep.detach()
+        }
+
+        return denoising_loss, critic_log_dict

predictor_training/__init__.py

@@ -0,0 +1,25 @@
+"""Offline Predictor-v4 training utilities."""
+
+from .dataset import (
+    SUPERVISION_PAIRS,
+    PredictorV4PairDataset,
+    move_batch_to_device,
+    predictor_v4_collate,
+)
+from .sampler import DistributedContextBucketBatchSampler
+from .trajectory_dataset import (
+    PredictorV4TrajectoryDataset,
+    load_offline_ffff_target,
+    trajectory_collate,
+)
+
+__all__ = [
+    "SUPERVISION_PAIRS",
+    "PredictorV4PairDataset",
+    "DistributedContextBucketBatchSampler",
+    "move_batch_to_device",
+    "predictor_v4_collate",
+    "PredictorV4TrajectoryDataset",
+    "load_offline_ffff_target",
+    "trajectory_collate",
+]

predictor_training/cache.py

@@ -0,0 +1,22 @@
+"""Small cache adapters for offline Predictor-v4 training."""
+
+from __future__ import annotations
+
+from typing import Mapping
+
+import torch
+
+
+def build_cross_attention_cache(
+    text_kv: Mapping[int, Mapping[str, torch.Tensor]],
+    source_block_ids: tuple[int, ...],
+) -> dict[int, dict[str, torch.Tensor | bool]]:
+    """Wrap saved Teacher text K/V in Wan's initialized-cache structure."""
+    return {
+        int(block_id): {
+            "k": text_kv[int(block_id)]["k"],
+            "v": text_kv[int(block_id)]["v"],
+            "is_init": True,
+        }
+        for block_id in source_block_ids
+    }

predictor_training/checkpoint.py

@@ -0,0 +1,103 @@
+"""Atomic Predictor-v4 inference and resumable training checkpoints."""
+
+from __future__ import annotations
+
+import json
+import os
+import random
+from pathlib import Path
+from typing import Any
+
+import torch
+from safetensors.torch import load_file, save_file
+
+
+def unwrap_model(model: torch.nn.Module) -> torch.nn.Module:
+    return model.module if hasattr(model, "module") else model
+
+
+def trainable_state_dict(
+    model: torch.nn.Module,
+    *,
+    floating_dtype: torch.dtype | None = None,
+) -> dict[str, torch.Tensor]:
+    model = unwrap_model(model)
+    trainable_names = {
+        name for name, parameter in model.named_parameters() if parameter.requires_grad
+    }
+    return {
+        name: tensor.detach()
+        .to(
+            device="cpu",
+            dtype=(
+                floating_dtype
+                if floating_dtype is not None and tensor.is_floating_point()
+                else tensor.dtype
+            ),
+        )
+        .contiguous()
+        for name, tensor in model.state_dict().items()
+        if name in trainable_names
+    }
+
+
+def save_predictor_weights(
+    model: torch.nn.Module,
+    path: str | Path,
+    *,
+    metadata: dict[str, Any],
+    floating_dtype: torch.dtype = torch.bfloat16,
+) -> Path:
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    temporary = path.with_suffix(path.suffix + f".tmp.{os.getpid()}")
+    save_file(
+        trainable_state_dict(model, floating_dtype=floating_dtype),
+        str(temporary),
+        metadata={
+            "format": "self_forcing_predictor_v4",
+            "config": json.dumps(metadata, ensure_ascii=False, sort_keys=True),
+        },
+    )
+    os.replace(temporary, path)
+    return path
+
+
+def load_predictor_weights(model: torch.nn.Module, path: str | Path) -> None:
+    state = load_file(str(path), device="cpu")
+    result = unwrap_model(model).load_state_dict(state, strict=False)
+    trainable = {
+        name
+        for name, parameter in unwrap_model(model).named_parameters()
+        if parameter.requires_grad
+    }
+    missing_trainable = sorted(trainable.intersection(result.missing_keys))
+    if result.unexpected_keys or missing_trainable:
+        raise RuntimeError(
+            "Predictor weight mismatch: "
+            f"unexpected={result.unexpected_keys}, "
+            f"missing_trainable={missing_trainable}"
+        )
+
+
+def atomic_torch_save(payload: dict[str, Any], path: str | Path) -> Path:
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    temporary = path.with_suffix(path.suffix + f".tmp.{os.getpid()}")
+    torch.save(payload, temporary)
+    os.replace(temporary, path)
+    return path
+
+
+def capture_rng_state() -> dict[str, Any]:
+    return {
+        "python": random.getstate(),
+        "torch_cpu": torch.get_rng_state(),
+        "torch_cuda": torch.cuda.get_rng_state(),
+    }
+
+
+def restore_rng_state(state: dict[str, Any]) -> None:
+    random.setstate(state["python"])
+    torch.set_rng_state(state["torch_cpu"])
+    torch.cuda.set_rng_state(state["torch_cuda"])

predictor_training/dataset.py

@@ -0,0 +1,364 @@
+"""Manifest-backed offline dataset for the three Predictor-v4 transitions.
+
+One manifest record represents one temporal chunk.  This dataset expands every
+usable record into the adjacent denoising pairs 0->1, 1->2 and 2->3.  Chunk
+zero is intentionally excluded because v4 conditions on the preceding chunk.
+"""
+
+from __future__ import annotations
+
+import json
+from functools import lru_cache
+from pathlib import Path
+from typing import Any, Iterable, Mapping
+
+import torch
+from safetensors import safe_open
+from torch.utils.data import Dataset
+
+
+SUPERVISION_PAIRS = ((0, 1), (1, 2), (2, 3))
+SCHEMA_VERSION = "self_forcing_predictor_v4_bf16_v1"
+FRAMES_PER_CHUNK = 3
+
+
+def _resolve(root: Path, value: str | Path) -> Path:
+    path = Path(value)
+    return path if path.is_absolute() else root / path
+
+
+def _load_selected(path: Path, names: Iterable[str]) -> dict[str, torch.Tensor]:
+    if not path.is_file():
+        raise FileNotFoundError(path)
+    with safe_open(str(path), framework="pt", device="cpu") as handle:
+        available = set(handle.keys())
+        missing = set(names).difference(available)
+        if missing:
+            raise KeyError(f"{path} is missing tensors {sorted(missing)}")
+        return {name: handle.get_tensor(name) for name in names}
+
+
+def _load_clean_prefeature(
+    path: Path,
+    candidates: Iterable[str],
+    *,
+    expected_start_frame: int,
+) -> torch.Tensor:
+    with safe_open(str(path), framework="pt", device="cpu") as handle:
+        available = set(handle.keys())
+        for name in candidates:
+            if name in available:
+                feature = handle.get_tensor(name)
+                break
+        else:
+            raise KeyError(
+                f"{path} has none of the expected tensors {tuple(candidates)}"
+            )
+        if "start_frame" in available:
+            actual_start = int(handle.get_tensor("start_frame").item())
+            if actual_start != expected_start_frame:
+                raise ValueError(
+                    f"{path} starts at frame {actual_start}, expected "
+                    f"{expected_start_frame}; history files are not ordered"
+                )
+        if "num_frames" in available:
+            actual_frames = int(handle.get_tensor("num_frames").item())
+            if actual_frames != FRAMES_PER_CHUNK:
+                raise ValueError(
+                    f"{path} contains {actual_frames} frames, expected "
+                    f"{FRAMES_PER_CHUNK}"
+                )
+        return feature
+
+
+def _block_entry(mapping: Mapping[Any, Any], block_id: int) -> Any:
+    for key in (str(block_id), block_id, f"block_{block_id}", f"block_{block_id:02d}"):
+        if key in mapping:
+            return mapping[key]
+    raise KeyError(f"No prefeature entry for block {block_id}")
+
+
+def _as_path_list(entry: Any) -> list[str]:
+    if isinstance(entry, (str, Path)):
+        return [str(entry)]
+    if isinstance(entry, Mapping):
+        # Builders may use {"files": [...]} or {"file": "..."}.
+        for key in ("files", "paths", "history", "file", "path"):
+            if key in entry:
+                return _as_path_list(entry[key])
+    if isinstance(entry, (list, tuple)):
+        return [str(value) for value in entry]
+    raise TypeError(f"Unsupported prefeature file entry: {entry!r}")
+
+
+def _prefeature_names(block_id: int) -> tuple[str, ...]:
+    return (
+        "self_attn_input",
+        "clean_prefeature",
+        "prefeature",
+        "img_modulated",
+        f"block_{block_id}_self_attn_input",
+        f"block_{block_id:02d}_self_attn_input",
+    )
+
+
+class PredictorV4PairDataset(Dataset):
+    """Read safetensors records and expose adjacent-step supervision pairs."""
+
+    PAIRS = SUPERVISION_PAIRS
+
+    def __init__(
+        self,
+        manifest_path: str | Path,
+        *,
+        source_block_ids: tuple[int, ...] = (1, 28),
+        max_records: int | None = None,
+        require_previous_chunk: bool = True,
+    ) -> None:
+        self.manifest_path = Path(manifest_path).resolve()
+        self.root = self.manifest_path.parent
+        self.source_block_ids = tuple(int(value) for value in source_block_ids)
+        with self.manifest_path.open("r", encoding="utf-8") as handle:
+            records = [json.loads(line) for line in handle if line.strip()]
+        if require_previous_chunk:
+            records = [record for record in records if int(record["chunk_id"]) > 0]
+        if max_records is not None:
+            records = records[: int(max_records)]
+        if not records:
+            raise ValueError(f"No usable records in {self.manifest_path}")
+        for record in records:
+            required = {
+                "step_tensor_file",
+                "previous_step_tensor_file",
+                "case_tensor_file",
+                "chunk_id",
+            }
+            missing = required.difference(record)
+            if missing:
+                raise ValueError(f"Manifest record lacks fields {sorted(missing)}")
+            if record.get("schema_version", SCHEMA_VERSION) != SCHEMA_VERSION:
+                raise ValueError(
+                    f"Unsupported Predictor schema {record.get('schema_version')!r}"
+                )
+            chunk_id = int(record["chunk_id"])
+            expected_context_frames = chunk_id * FRAMES_PER_CHUNK
+            context_frames = int(
+                record.get("context_frames", expected_context_frames)
+            )
+            if context_frames != expected_context_frames:
+                raise ValueError(
+                    f"chunk {chunk_id} context_frames={context_frames}, expected "
+                    f"{expected_context_frames}"
+                )
+            history = record.get("history_clean_prefeature_files")
+            if history is None:
+                raise ValueError(
+                    "Manifest record lacks history_clean_prefeature_files; "
+                    "clean_prefeature_files contains only the current chunk"
+                )
+            for block_id in self.source_block_ids:
+                paths = _as_path_list(_block_entry(history, block_id))
+                if len(paths) != int(record["chunk_id"]):
+                    raise ValueError(
+                        f"chunk {record['chunk_id']} block {block_id} has "
+                        f"{len(paths)} history files, expected {record['chunk_id']}"
+                    )
+        self.records = records
+
+    def __len__(self) -> int:
+        return len(self.records) * len(self.PAIRS)
+
+    @lru_cache(maxsize=8)
+    def _load_case(self, relative_path: str) -> dict[str, torch.Tensor]:
+        path = _resolve(self.root, relative_path)
+        names: list[str] = []
+        with safe_open(str(path), framework="pt", device="cpu") as handle:
+            keys = set(handle.keys())
+            for block_id in self.source_block_ids:
+                for kind in ("k", "v"):
+                    candidates = (
+                        f"block_{block_id:02d}_cross_{kind}",
+                        f"block_{block_id}_cross_{kind}",
+                        f"block_{block_id}_text_{kind}",
+                        f"block_{block_id:02d}_text_{kind}",
+                        f"block_{block_id}_{kind}_txt",
+                        f"text_{kind}_block_{block_id}",
+                    )
+                    found = next((name for name in candidates if name in keys), None)
+                    if found is None:
+                        raise KeyError(
+                            f"{path} has no text {kind.upper()} for block {block_id}"
+                        )
+                    names.append(found)
+            return {name: handle.get_tensor(name) for name in names}
+
+    def _case_text_kv(
+        self,
+        relative_path: str,
+    ) -> dict[int, dict[str, torch.Tensor]]:
+        tensors = self._load_case(relative_path)
+        result: dict[int, dict[str, torch.Tensor]] = {}
+        for block_id in self.source_block_ids:
+            result[block_id] = {}
+            for kind in ("k", "v"):
+                candidates = (
+                    f"block_{block_id:02d}_cross_{kind}",
+                    f"block_{block_id}_cross_{kind}",
+                    f"block_{block_id}_text_{kind}",
+                    f"block_{block_id:02d}_text_{kind}",
+                    f"block_{block_id}_{kind}_txt",
+                    f"text_{kind}_block_{block_id}",
+                )
+                name = next(name for name in candidates if name in tensors)
+                result[block_id][kind] = tensors[name]
+        return result
+
+    def _history_prefeature(
+        self,
+        record: Mapping[str, Any],
+    ) -> dict[int, torch.Tensor]:
+        history = record["history_clean_prefeature_files"]
+        result = {}
+        for block_id in self.source_block_ids:
+            paths = _as_path_list(_block_entry(history, block_id))
+            chunks = [
+                _load_clean_prefeature(
+                    _resolve(self.root, path),
+                    _prefeature_names(block_id),
+                    expected_start_frame=chunk_index * FRAMES_PER_CHUNK,
+                )
+                for chunk_index, path in enumerate(paths)
+            ]
+            # Files are [1, chunk_tokens, dim].  Concatenate temporal chunks.
+            result[block_id] = torch.cat(chunks, dim=1)
+        return result
+
+    def __getitem__(self, index: int) -> dict[str, Any]:
+        record_index, pair_index = divmod(index, len(self.PAIRS))
+        record = self.records[record_index]
+        anchor_step, target_step = self.PAIRS[pair_index]
+        step_path = _resolve(self.root, record["step_tensor_file"])
+        step_names = (
+            f"step_{anchor_step}_final_hidden",
+            f"step_{target_step}_noisy_latent",
+            f"step_{target_step}_timestep",
+            f"step_{target_step}_final_hidden",
+            f"step_{target_step}_flow",
+        )
+        step_tensors = _load_selected(step_path, step_names)
+        previous_name = f"step_{target_step}_final_hidden"
+        previous = _load_selected(
+            _resolve(self.root, record["previous_step_tensor_file"]),
+            (previous_name,),
+        )
+        context_frames = int(
+            record.get(
+                "context_frames",
+                int(record["chunk_id"]) * FRAMES_PER_CHUNK,
+            )
+        )
+        return {
+            "target_latent": step_tensors[f"step_{target_step}_noisy_latent"],
+            "target_timestep": step_tensors[f"step_{target_step}_timestep"],
+            "anchor_hidden": step_tensors[f"step_{anchor_step}_final_hidden"],
+            "previous_chunk_hidden": previous[previous_name],
+            "target_hidden": step_tensors[f"step_{target_step}_final_hidden"],
+            "target_flow": step_tensors[f"step_{target_step}_flow"],
+            "clean_prefeature": self._history_prefeature(record),
+            "text_kv": self._case_text_kv(str(record["case_tensor_file"])),
+            "case_id": record.get("case_id"),
+            "chunk_id": int(record["chunk_id"]),
+            "context_frames": context_frames,
+            "anchor_step": anchor_step,
+            "target_step": target_step,
+        }
+
+
+def predictor_v4_collate(items: list[dict[str, Any]]) -> dict[str, Any]:
+    """Collate a context-length bucket into one batch."""
+    if not items:
+        raise ValueError("Cannot collate an empty batch")
+    context_frames = {item["context_frames"] for item in items}
+    if len(context_frames) != 1:
+        raise ValueError(
+            "A batch must have one history length; enable bucket_by_context"
+        )
+    tensor_keys = (
+        "target_latent",
+        "target_timestep",
+        "anchor_hidden",
+        "previous_chunk_hidden",
+        "target_hidden",
+        "target_flow",
+    )
+    batch: dict[str, Any] = {
+        key: torch.cat([item[key] for item in items], dim=0) for key in tensor_keys
+    }
+    block_ids = tuple(items[0]["clean_prefeature"])
+    batch["clean_prefeature"] = {
+        block_id: torch.cat(
+            [item["clean_prefeature"][block_id] for item in items], dim=0
+        )
+        for block_id in block_ids
+    }
+    batch["text_kv"] = {
+        block_id: {
+            kind: torch.cat(
+                [item["text_kv"][block_id][kind] for item in items], dim=0
+            )
+            for kind in ("k", "v")
+        }
+        for block_id in block_ids
+    }
+    for key in (
+        "case_id",
+        "chunk_id",
+        "context_frames",
+        "anchor_step",
+        "target_step",
+    ):
+        batch[key] = [item[key] for item in items]
+    return batch
+
+
+def _move(
+    tensor: torch.Tensor,
+    *,
+    device: torch.device,
+    dtype: torch.dtype,
+) -> torch.Tensor:
+    target_dtype = dtype if tensor.is_floating_point() else tensor.dtype
+    return tensor.to(device=device, dtype=target_dtype, non_blocking=True)
+
+
+def move_batch_to_device(
+    batch: dict[str, Any],
+    *,
+    device: torch.device,
+    dtype: torch.dtype,
+) -> dict[str, Any]:
+    result = {
+        key: _move(batch[key], device=device, dtype=dtype)
+        for key in (
+            "target_latent",
+            "target_timestep",
+            "anchor_hidden",
+            "previous_chunk_hidden",
+            "target_hidden",
+            "target_flow",
+        )
+    }
+    result["clean_prefeature"] = {
+        int(block_id): _move(value, device=device, dtype=dtype)
+        for block_id, value in batch["clean_prefeature"].items()
+    }
+    result["text_kv"] = {
+        int(block_id): {
+            kind: _move(value, device=device, dtype=dtype)
+            for kind, value in values.items()
+        }
+        for block_id, values in batch["text_kv"].items()
+    }
+    for key in ("case_id", "chunk_id", "context_frames", "anchor_step", "target_step"):
+        result[key] = batch[key]
+    return result

predictor_training/rollout_cache.py

@@ -0,0 +1,114 @@
+"""Cache workspaces for detached Predictor-v4 trajectory rollout."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+import torch
+
+
+def reset_main_caches(
+    kv_cache: Sequence[dict],
+    crossattn_cache: Sequence[dict],
+) -> None:
+    """Reset logical cache extents without clearing unused storage."""
+
+    for cache in kv_cache:
+        cache["global_end_index"].zero_()
+        cache["local_end_index"].zero_()
+    for cache in crossattn_cache:
+        cache["is_init"] = False
+
+
+def build_predictor_workspace(
+    main_cache: Sequence[dict],
+    *,
+    source_block_ids: tuple[int, int],
+    history_tokens: int,
+    current_tokens: int,
+) -> dict[int, dict[str, torch.Tensor]]:
+    """Copy committed clean history into a temporary selected-layer workspace.
+
+    The current-chunk region is intentionally empty. Predictor P1/P2/P3
+    overwrite that region, and the entire workspace is discarded after the
+    chunk. Only the final Full timestep-zero pass updates persistent history.
+    """
+
+    history_tokens = int(history_tokens)
+    current_tokens = int(current_tokens)
+    if history_tokens < 0 or current_tokens <= 0:
+        raise ValueError("Invalid history/current token count")
+    capacity = history_tokens + current_tokens
+    result: dict[int, dict[str, torch.Tensor]] = {}
+    for block_id in source_block_ids:
+        source = main_cache[int(block_id)]
+        if int(source["global_end_index"].item()) < history_tokens:
+            raise RuntimeError(
+                f"Teacher cache block {block_id} ends before committed history: "
+                f"{int(source['global_end_index'].item())} < {history_tokens}"
+            )
+        key = source["k"].new_zeros(
+            source["k"].shape[0], capacity, *source["k"].shape[2:]
+        )
+        value = source["v"].new_zeros(
+            source["v"].shape[0], capacity, *source["v"].shape[2:]
+        )
+        if history_tokens:
+            key[:, :history_tokens].copy_(source["k"][:, :history_tokens])
+            value[:, :history_tokens].copy_(source["v"][:, :history_tokens])
+        result[int(block_id)] = {
+            "k": key,
+            "v": value,
+            "global_end_index": torch.tensor(
+                [history_tokens], dtype=torch.long, device=key.device
+            ),
+            "local_end_index": torch.tensor(
+                [history_tokens], dtype=torch.long, device=key.device
+            ),
+        }
+    return result
+
+
+def reset_predictor_workspace(
+    workspace: dict[int, dict[str, torch.Tensor]],
+    *,
+    history_tokens: int,
+) -> None:
+    """Discard the previous timestep's differentiable current-cache region."""
+
+    history_tokens = int(history_tokens)
+    for cache in workspace.values():
+        # Slice assignment from Predictor K/V can attach CopySlices autograd
+        # history to the workspace tensor. Replace it with a detached tensor
+        # before the next timestep so P2/P3 cannot backpropagate through cache.
+        cache["k"] = cache["k"].detach()
+        cache["v"] = cache["v"].detach()
+        if history_tokens < cache["k"].shape[1]:
+            cache["k"][:, history_tokens:].zero_()
+            cache["v"][:, history_tokens:].zero_()
+        cache["global_end_index"].fill_(history_tokens)
+        cache["local_end_index"].fill_(history_tokens)
+
+
+def assert_clean_history_extent(
+    kv_cache: Sequence[dict],
+    *,
+    expected_tokens: int,
+) -> None:
+    expected_tokens = int(expected_tokens)
+    for block_id, cache in enumerate(kv_cache):
+        global_end = int(cache["global_end_index"].item())
+        local_end = int(cache["local_end_index"].item())
+        if global_end != expected_tokens or local_end != expected_tokens:
+            raise RuntimeError(
+                f"Cache block {block_id} extent is ({global_end}, {local_end}), "
+                f"expected committed clean history {expected_tokens}"
+            )
+
+
+__all__ = [
+    "assert_clean_history_extent",
+    "build_predictor_workspace",
+    "reset_main_caches",
+    "reset_predictor_workspace",
+]

predictor_training/sampler.py

@@ -0,0 +1,84 @@
+"""Deterministic DDP batches grouped by clean-history length."""
+
+from __future__ import annotations
+
+import math
+import random
+from collections import defaultdict
+from typing import Iterator
+
+from torch.utils.data import Sampler
+
+from .dataset import FRAMES_PER_CHUNK
+
+
+class DistributedContextBucketBatchSampler(Sampler[list[int]]):
+    def __init__(
+        self,
+        dataset,
+        *,
+        batch_size: int,
+        rank: int,
+        world_size: int,
+        seed: int = 0,
+        drop_last: bool = True,
+    ) -> None:
+        self.dataset = dataset
+        self.batch_size = int(batch_size)
+        self.rank = int(rank)
+        self.world_size = int(world_size)
+        self.seed = int(seed)
+        self.drop_last = bool(drop_last)
+        self.epoch = 0
+        if self.batch_size <= 0 or not 0 <= self.rank < self.world_size:
+            raise ValueError("Invalid distributed bucket sampler configuration")
+        self._buckets: dict[int, list[int]] = defaultdict(list)
+        for record_index, record in enumerate(dataset.records):
+            context_frames = int(
+                record.get(
+                    "context_frames",
+                    int(record["chunk_id"]) * FRAMES_PER_CHUNK,
+                )
+            )
+            for pair_index in range(len(dataset.PAIRS)):
+                self._buckets[context_frames].append(
+                    record_index * len(dataset.PAIRS) + pair_index
+                )
+
+    def set_epoch(self, epoch: int) -> None:
+        self.epoch = int(epoch)
+
+    def _global_batches(self) -> list[list[int]]:
+        rng = random.Random(self.seed + self.epoch)
+        global_batch_size = self.batch_size * self.world_size
+        batches: list[list[int]] = []
+        for bucket in self._buckets.values():
+            indices = list(bucket)
+            rng.shuffle(indices)
+            if not self.drop_last and len(indices) % global_batch_size:
+                needed = global_batch_size - len(indices) % global_batch_size
+                indices.extend((indices * math.ceil(needed / len(indices)))[:needed])
+            usable = len(indices) - len(indices) % global_batch_size
+            batches.extend(
+                indices[start : start + global_batch_size]
+                for start in range(0, usable, global_batch_size)
+            )
+        rng.shuffle(batches)
+        return batches
+
+    def __iter__(self) -> Iterator[list[int]]:
+        start = self.rank * self.batch_size
+        end = start + self.batch_size
+        for global_batch in self._global_batches():
+            yield global_batch[start:end]
+
+    def __len__(self) -> int:
+        global_batch_size = self.batch_size * self.world_size
+        if self.drop_last:
+            return sum(
+                len(values) // global_batch_size for values in self._buckets.values()
+            )
+        return sum(
+            math.ceil(len(values) / global_batch_size)
+            for values in self._buckets.values()
+        )

predictor_training/trajectory_dataset.py

@@ -0,0 +1,107 @@
+"""Prompt trajectories and fixed FFFF targets for Predictor-v4 Stage 2."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any
+
+import torch
+from safetensors import safe_open
+from torch.utils.data import Dataset
+
+
+class PredictorV4TrajectoryDataset(Dataset):
+    """One deterministic seven-chunk trajectory per training prompt."""
+
+    def __init__(
+        self,
+        cases_path: str | Path,
+        *,
+        data_root: str | Path,
+        max_cases: int | None = None,
+    ) -> None:
+        self.cases_path = Path(cases_path).resolve()
+        self.data_root = Path(data_root).resolve()
+        with self.cases_path.open("r", encoding="utf-8") as handle:
+            cases = [json.loads(line) for line in handle if line.strip()]
+        if max_cases is not None:
+            cases = cases[: int(max_cases)]
+        if not cases:
+            raise ValueError(f"No trajectory cases in {self.cases_path}")
+        for position, case in enumerate(cases):
+            if int(case["case_id"]) != position:
+                raise ValueError("Stage-2 case_id values must be dense and ordered")
+            if int(case.get("seed", -1)) != 0:
+                raise ValueError(f"case {position} does not use latent seed 0")
+            if not str(case.get("prompt", "")).strip():
+                raise ValueError(f"case {position} has an empty prompt")
+        self.cases = cases
+
+    def __len__(self) -> int:
+        return len(self.cases)
+
+    def __getitem__(self, index: int) -> dict[str, Any]:
+        case = self.cases[index]
+        return {
+            "case_id": int(case["case_id"]),
+            "prompt": str(case["prompt"]),
+            "seed": int(case["seed"]),
+        }
+
+    def offline_step_path(self, case_id: int, chunk_id: int) -> Path:
+        path = (
+            self.data_root
+            / "steps"
+            / f"case_{int(case_id):06d}"
+            / f"chunk_{int(chunk_id):02d}.safetensors"
+        )
+        if not path.is_file():
+            raise FileNotFoundError(path)
+        return path
+
+
+def trajectory_collate(items: list[dict[str, Any]]) -> dict[str, Any]:
+    if len(items) != 1:
+        raise ValueError(
+            "Stage-2 rollout currently requires one trajectory per rank"
+        )
+    return items[0]
+
+
+def load_offline_ffff_target(
+    path: str | Path,
+    *,
+    step_id: int,
+    device: torch.device,
+) -> dict[str, torch.Tensor]:
+    """Load one immutable Full-trajectory hidden/flow target."""
+
+    path = Path(path)
+    names = {
+        "hidden": f"step_{int(step_id)}_final_hidden",
+        "flow": f"step_{int(step_id)}_flow",
+        "timestep": f"step_{int(step_id)}_timestep",
+    }
+    with safe_open(str(path), framework="pt", device="cpu") as handle:
+        missing = set(names.values()).difference(handle.keys())
+        if missing:
+            raise KeyError(f"{path} lacks fixed FFFF tensors {sorted(missing)}")
+        result = {key: handle.get_tensor(name) for key, name in names.items()}
+    if result["hidden"].dtype != torch.bfloat16:
+        raise TypeError(f"{path}: target hidden must be BF16")
+    if result["flow"].dtype != torch.bfloat16:
+        raise TypeError(f"{path}: target flow must be BF16")
+    if result["timestep"].dtype != torch.int64:
+        raise TypeError(f"{path}: legacy offline timestep must be INT64")
+    return {
+        key: value.to(device=device, non_blocking=True)
+        for key, value in result.items()
+    }
+
+
+__all__ = [
+    "PredictorV4TrajectoryDataset",
+    "load_offline_ffff_target",
+    "trajectory_collate",
+]

utils/lmdb.py

@@ -0,0 +1,72 @@
+import numpy as np
+
+
+def get_array_shape_from_lmdb(env, array_name):
+    with env.begin() as txn:
+        image_shape = txn.get(f"{array_name}_shape".encode()).decode()
+        image_shape = tuple(map(int, image_shape.split()))
+    return image_shape
+
+
+def store_arrays_to_lmdb(env, arrays_dict, start_index=0):
+    """
+    Store rows of multiple numpy arrays in a single LMDB.
+    Each row is stored separately with a naming convention.
+    """
+    with env.begin(write=True) as txn:
+        for array_name, array in arrays_dict.items():
+            for i, row in enumerate(array):
+                # Convert row to bytes
+                if isinstance(row, str):
+                    row_bytes = row.encode()
+                else:
+                    row_bytes = row.tobytes()
+
+                data_key = f'{array_name}_{start_index + i}_data'.encode()
+
+                txn.put(data_key, row_bytes)
+
+
+def process_data_dict(data_dict, seen_prompts):
+    output_dict = {}
+
+    all_videos = []
+    all_prompts = []
+    for prompt, video in data_dict.items():
+        if prompt in seen_prompts:
+            continue
+        else:
+            seen_prompts.add(prompt)
+
+        video = video.half().numpy()
+        all_videos.append(video)
+        all_prompts.append(prompt)
+
+    if len(all_videos) == 0:
+        return {"latents": np.array([]), "prompts": np.array([])}
+
+    all_videos = np.concatenate(all_videos, axis=0)
+
+    output_dict['latents'] = all_videos
+    output_dict['prompts'] = np.array(all_prompts)
+
+    return output_dict
+
+
+def retrieve_row_from_lmdb(lmdb_env, array_name, dtype, row_index, shape=None):
+    """
+    Retrieve a specific row from a specific array in the LMDB.
+    """
+    data_key = f'{array_name}_{row_index}_data'.encode()
+
+    with lmdb_env.begin() as txn:
+        row_bytes = txn.get(data_key)
+
+    if dtype == str:
+        array = row_bytes.decode()
+    else:
+        array = np.frombuffer(row_bytes, dtype=dtype)
+
+    if shape is not None and len(shape) > 0:
+        array = array.reshape(shape)
+    return array

wan_models/Wan2.1-T2V-1.3B/.gitattributes

@@ -0,0 +1,47 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+google/umt5-xxl/tokenizer.json filter=lfs diff=lfs merge=lfs -text
+assets/comp_effic.png filter=lfs diff=lfs merge=lfs -text
+assets/data_for_diff_stage.jpg filter=lfs diff=lfs merge=lfs -text
+assets/i2v_res.png filter=lfs diff=lfs merge=lfs -text
+assets/logo.png filter=lfs diff=lfs merge=lfs -text
+assets/t2v_res.jpg filter=lfs diff=lfs merge=lfs -text
+assets/vben_vs_sota.png filter=lfs diff=lfs merge=lfs -text
+assets/vben_vs_sota_t2i.jpg filter=lfs diff=lfs merge=lfs -text
+assets/video_dit_arch.jpg filter=lfs diff=lfs merge=lfs -text
+assets/video_vae_res.jpg filter=lfs diff=lfs merge=lfs -text
+examples/i2v_input.JPG filter=lfs diff=lfs merge=lfs -text
+assets/.DS_Store filter=lfs diff=lfs merge=lfs -text

wan_models/Wan2.1-T2V-1.3B/LICENSE.txt

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   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.

wan_models/Wan2.1-T2V-1.3B/README.md

@@ -0,0 +1,298 @@
+---
+license: apache-2.0
+language:
+- en
+- zh
+pipeline_tag: text-to-video
+library_name: diffusers
+tags:
+- video
+- video-generation
+---
+# Wan2.1
+
+<p align="center">
+    <img src="assets/logo.png" width="400"/>
+<p>
+
+<p align="center">
+    💜 <a href=""><b>Wan</b></a> &nbsp&nbsp ｜ &nbsp&nbsp 🖥️ <a href="https://github.com/Wan-Video/Wan2.1">GitHub</a> &nbsp&nbsp  | &nbsp&nbsp🤗 <a href="https://huggingface.co/Wan-AI/">Hugging Face</a>&nbsp&nbsp | &nbsp&nbsp🤖 <a href="https://modelscope.cn/organization/Wan-AI">ModelScope</a>&nbsp&nbsp | &nbsp&nbsp 📑 <a href="">Paper (Coming soon)</a> &nbsp&nbsp | &nbsp&nbsp 📑 <a href="https://wanxai.com">Blog</a> &nbsp&nbsp | &nbsp&nbsp💬 <a href="https://gw.alicdn.com/imgextra/i2/O1CN01tqjWFi1ByuyehkTSB_!!6000000000015-0-tps-611-1279.jpg">WeChat Group</a>&nbsp&nbsp | &nbsp&nbsp 📖 <a href="https://discord.gg/p5XbdQV7">Discord</a>&nbsp&nbsp
+<br>
+
+-----
+
+[**Wan: Open and Advanced Large-Scale Video Generative Models**]("#") <be>
+
+In this repository, we present **Wan2.1**, a comprehensive and open suite of video foundation models that pushes the boundaries of video generation. **Wan2.1** offers these key features:
+- 👍 **SOTA Performance**: **Wan2.1** consistently outperforms existing open-source models and state-of-the-art commercial solutions across multiple benchmarks.
+- 👍 **Supports Consumer-grade GPUs**: The T2V-1.3B model requires only 8.19 GB VRAM, making it compatible with almost all consumer-grade GPUs. It can generate a 5-second 480P video on an RTX 4090 in about 4 minutes (without optimization techniques like quantization). Its performance is even comparable to some closed-source models.
+- 👍 **Multiple Tasks**: **Wan2.1** excels in Text-to-Video, Image-to-Video, Video Editing, Text-to-Image, and Video-to-Audio, advancing the field of video generation.
+- 👍 **Visual Text Generation**: **Wan2.1** is the first video model capable of generating both Chinese and English text, featuring robust text generation that enhances its practical applications.
+- 👍 **Powerful Video VAE**: **Wan-VAE** delivers exceptional efficiency and performance, encoding and decoding 1080P videos of any length while preserving temporal information, making it an ideal foundation for video and image generation.
+
+
+This repository hosts our T2V-1.3B model, a versatile solution for video generation that is compatible with nearly all consumer-grade GPUs. In this way, we hope that **Wan2.1** can serve as an easy-to-use tool for more creative teams in video creation, providing a high-quality foundational model for academic teams with limited computing resources. This will facilitate both the rapid development of the video creation community and the swift advancement of video technology.
+
+
+## Video Demos
+
+<div align="center">
+    <video width="80%" controls>
+        <source src="https://cloud.video.taobao.com/vod/Jth64Y7wNoPcJki_Bo1ZJTDBvNjsgjlVKsNs05Fqfps.mp4" type="video/mp4">
+        Your browser does not support the video tag.
+    </video>
+</div>
+
+
+## 🔥 Latest News!!
+
+* Feb 25, 2025: 👋 We've released the inference code and weights of Wan2.1.
+
+
+## 📑 Todo List
+- Wan2.1 Text-to-Video
+    - [x] Multi-GPU Inference code of the 14B and 1.3B models
+    - [x] Checkpoints of the 14B and 1.3B models
+    - [x] Gradio demo
+    - [ ] Diffusers integration
+    - [ ] ComfyUI integration
+- Wan2.1 Image-to-Video
+    - [x] Multi-GPU Inference code of the 14B model
+    - [x] Checkpoints of the 14B model
+    - [x] Gradio demo
+    - [ ] Diffusers integration
+    - [ ] ComfyUI integration
+
+
+## Quickstart
+
+#### Installation
+Clone the repo:
+```
+git clone https://github.com/Wan-Video/Wan2.1.git
+cd Wan2.1
+```
+
+Install dependencies:
+```
+# Ensure torch >= 2.4.0
+pip install -r requirements.txt
+```
+
+
+#### Model Download
+
+| Models        |                       Download Link                                           |    Notes                      |
+| --------------|-------------------------------------------------------------------------------|-------------------------------|
+| T2V-14B       |      🤗 [Huggingface](https://huggingface.co/Wan-AI/Wan2.1-T2V-14B)      🤖 [ModelScope](https://www.modelscope.cn/models/Wan-AI/Wan2.1-T2V-14B)          | Supports both 480P and 720P
+| I2V-14B-720P  |      🤗 [Huggingface](https://huggingface.co/Wan-AI/Wan2.1-I2V-14B-720P)    🤖 [ModelScope](https://www.modelscope.cn/models/Wan-AI/Wan2.1-I2V-14B-720P)     | Supports 720P
+| I2V-14B-480P  |      🤗 [Huggingface](https://huggingface.co/Wan-AI/Wan2.1-I2V-14B-480P)    🤖 [ModelScope](https://www.modelscope.cn/models/Wan-AI/Wan2.1-I2V-14B-480P)      | Supports 480P
+| T2V-1.3B      |      🤗 [Huggingface](https://huggingface.co/Wan-AI/Wan2.1-T2V-1.3B)     🤖 [ModelScope](https://www.modelscope.cn/models/Wan-AI/Wan2.1-T2V-1.3B)         | Supports 480P
+
+
+> 💡Note: The 1.3B model is capable of generating videos at 720P resolution. However, due to limited training at this resolution, the results are generally less stable compared to 480P. For optimal performance, we recommend using 480P resolution.
+
+
+Download models using 🤗 huggingface-cli:
+```
+pip install "huggingface_hub[cli]"
+huggingface-cli download Wan-AI/Wan2.1-T2V-1.3B --local-dir ./Wan2.1-T2V-1.3B
+```
+
+Download models using 🤖 modelscope-cli:
+```
+pip install modelscope
+modelscope download Wan-AI/Wan2.1-T2V-1.3B --local_dir ./Wan2.1-T2V-1.3B
+```
+
+#### Run Text-to-Video Generation
+
+This repository supports two Text-to-Video models (1.3B and 14B) and two resolutions (480P and 720P). The parameters and configurations for these models are as follows:
+
+<table>
+    <thead>
+        <tr>
+            <th rowspan="2">Task</th>
+            <th colspan="2">Resolution</th>
+            <th rowspan="2">Model</th>
+        </tr>
+        <tr>
+            <th>480P</th>
+            <th>720P</th>
+        </tr>
+    </thead>
+    <tbody>
+        <tr>
+            <td>t2v-14B</td>
+            <td style="color: green;">✔️</td>
+            <td style="color: green;">✔️</td>
+            <td>Wan2.1-T2V-14B</td>
+        </tr>
+        <tr>
+            <td>t2v-1.3B</td>
+            <td style="color: green;">✔️</td>
+            <td style="color: red;">❌</td>
+            <td>Wan2.1-T2V-1.3B</td>
+        </tr>
+    </tbody>
+</table>
+
+
+##### (1) Without Prompt Extention
+
+To facilitate implementation, we will start with a basic version of the inference process that skips the [prompt extension](#2-using-prompt-extention) step.
+
+- Single-GPU inference
+
+```
+python generate.py  --task t2v-1.3B --size 832*480 --ckpt_dir ./Wan2.1-T2V-1.3B --sample_shift 8 --sample_guide_scale 6 --prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage."
+```
+
+If you encounter OOM (Out-of-Memory) issues, you can use the `--offload_model True` and `--t5_cpu` options to reduce GPU memory usage. For example, on an RTX 4090 GPU:
+
+```
+python generate.py  --task t2v-1.3B --size 832*480 --ckpt_dir ./Wan2.1-T2V-1.3B --offload_model True --t5_cpu --sample_shift 8 --sample_guide_scale 6 --prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage."
+```
+
+> 💡Note: If you are using the `T2V-1.3B` model, we recommend setting the parameter `--sample_guide_scale 6`. The `--sample_shift parameter` can be adjusted within the range of 8 to 12 based on the performance.
+
+- Multi-GPU inference using FSDP + xDiT USP
+
+```
+pip install "xfuser>=0.4.1"
+torchrun --nproc_per_node=8 generate.py --task t2v-1.3B --size 832*480 --ckpt_dir ./Wan2.1-T2V-1.3B --dit_fsdp --t5_fsdp --ulysses_size 8 --sample_shift 8 --sample_guide_scale 6 --prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage."
+```
+
+
+##### (2) Using Prompt Extention
+
+Extending the prompts can effectively enrich the details in the generated videos, further enhancing the video quality. Therefore, we recommend enabling prompt extension. We provide the following two methods for prompt extension:
+
+- Use the Dashscope API for extension.
+  - Apply for a `dashscope.api_key` in advance ([EN](https://www.alibabacloud.com/help/en/model-studio/getting-started/first-api-call-to-qwen) | [CN](https://help.aliyun.com/zh/model-studio/getting-started/first-api-call-to-qwen)).
+  - Configure the environment variable `DASH_API_KEY` to specify the Dashscope API key. For users of Alibaba Cloud's international site, you also need to set the environment variable `DASH_API_URL` to 'https://dashscope-intl.aliyuncs.com/api/v1'. For more detailed instructions, please refer to the [dashscope document](https://www.alibabacloud.com/help/en/model-studio/developer-reference/use-qwen-by-calling-api?spm=a2c63.p38356.0.i1).
+  - Use the `qwen-plus` model for text-to-video tasks and `qwen-vl-max` for image-to-video tasks.
+  - You can modify the model used for extension with the parameter `--prompt_extend_model`. For example:
+```
+DASH_API_KEY=your_key python generate.py  --task t2v-1.3B --size 832*480 --ckpt_dir ./Wan2.1-T2V-1.3B --prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage" --use_prompt_extend --prompt_extend_method 'dashscope' --prompt_extend_target_lang 'ch'
+```
+
+- Using a local model for extension.
+
+  - By default, the Qwen model on HuggingFace is used for this extension. Users can choose based on the available GPU memory size.
+  - For text-to-video tasks, you can use models like `Qwen/Qwen2.5-14B-Instruct`, `Qwen/Qwen2.5-7B-Instruct` and `Qwen/Qwen2.5-3B-Instruct`
+  - For image-to-video tasks, you can use models like `Qwen/Qwen2.5-VL-7B-Instruct` and `Qwen/Qwen2.5-VL-3B-Instruct`.
+  - Larger models generally provide better extension results but require more GPU memory.
+  - You can modify the model used for extension with the parameter `--prompt_extend_model` , allowing you to specify either a local model path or a Hugging Face model. For example:
+
+```
+python generate.py  --task t2v-1.3B --size 832*480 --ckpt_dir ./Wan2.1-T2V-1.3B --prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage" --use_prompt_extend --prompt_extend_method 'local_qwen' --prompt_extend_target_lang 'ch'
+```
+
+##### (3) Runing local gradio
+
+```
+cd gradio
+# if one uses dashscope’s API for prompt extension
+DASH_API_KEY=your_key python t2v_1.3B_singleGPU.py --prompt_extend_method 'dashscope' --ckpt_dir ./Wan2.1-T2V-1.3B
+
+# if one uses a local model for prompt extension
+python t2v_1.3B_singleGPU.py --prompt_extend_method 'local_qwen' --ckpt_dir ./Wan2.1-T2V-1.3B
+```
+
+
+
+## Evaluation
+
+We employ our **Wan-Bench** framework to evaluate the performance of the T2V-1.3B model, with the results displayed in the table below. The results indicate that our smaller 1.3B model surpasses the overall metrics of larger open-source models, demonstrating the effectiveness of **WanX2.1**'s architecture and the data construction pipeline.
+
+<div align="center">
+    <img src="assets/vben_1.3b_vs_sota.png" alt="" style="width: 80%;" />
+</div>
+
+
+
+## Computational Efficiency on Different GPUs
+
+We test the computational efficiency of different **Wan2.1** models on different GPUs in the following table. The results are presented in the format: **Total time (s) / peak GPU memory (GB)**.
+
+
+<div align="center">
+    <img src="assets/comp_effic.png" alt="" style="width: 80%;" />
+</div>
+
+> The parameter settings for the tests presented in this table are as follows:
+> (1) For the 1.3B model on 8 GPUs, set `--ring_size 8` and `--ulysses_size 1`;
+> (2) For the 14B model on 1 GPU, use `--offload_model True`;
+> (3) For the 1.3B model on a single 4090 GPU, set `--offload_model True --t5_cpu`;
+> (4) For all testings, no prompt extension was applied, meaning `--use_prompt_extend` was not enabled.
+
+-------
+
+## Introduction of Wan2.1
+
+**Wan2.1**  is designed on the mainstream diffusion transformer paradigm, achieving significant advancements in generative capabilities through a series of innovations. These include our novel spatio-temporal variational autoencoder (VAE), scalable training strategies, large-scale data construction, and automated evaluation metrics. Collectively, these contributions enhance the model’s performance and versatility.
+
+
+##### (1) 3D Variational Autoencoders
+We propose a novel 3D causal VAE architecture, termed **Wan-VAE** specifically designed for video generation. By combining multiple strategies, we improve spatio-temporal compression, reduce memory usage, and ensure temporal causality. **Wan-VAE** demonstrates significant advantages in performance efficiency compared to other open-source VAEs. Furthermore, our **Wan-VAE** can encode and decode unlimited-length 1080P videos without losing historical temporal information, making it particularly well-suited for video generation tasks.
+
+
+<div align="center">
+    <img src="assets/video_vae_res.jpg" alt="" style="width: 80%;" />
+</div>
+
+
+##### (2) Video Diffusion DiT
+
+**Wan2.1** is designed using the Flow Matching framework within the paradigm of mainstream Diffusion Transformers. Our model's architecture uses the T5 Encoder to encode multilingual text input, with cross-attention in each transformer block embedding the text into the model structure. Additionally, we employ an MLP with a Linear layer and a SiLU layer to process the input time embeddings and predict six modulation parameters individually. This MLP is shared across all transformer blocks, with each block learning a distinct set of biases. Our experimental findings reveal a significant performance improvement with this approach at the same parameter scale.
+
+<div align="center">
+    <img src="assets/video_dit_arch.jpg" alt="" style="width: 80%;" />
+</div>
+
+
+| Model  | Dimension | Input Dimension | Output Dimension | Feedforward Dimension | Frequency Dimension | Number of Heads | Number of Layers |
+|--------|-----------|-----------------|------------------|-----------------------|---------------------|-----------------|------------------|
+| 1.3B   | 1536      | 16              | 16               | 8960                  | 256                 | 12              | 30               |
+| 14B   | 5120       | 16              | 16               | 13824                 | 256                 | 40              | 40               |
+
+
+
+##### Data
+
+We curated and deduplicated a candidate dataset comprising a vast amount of image and video data. During the data curation process, we designed a four-step data cleaning process, focusing on fundamental dimensions, visual quality and motion quality. Through the robust data processing pipeline, we can easily obtain high-quality, diverse, and large-scale training sets of images and videos.
+
+![figure1](assets/data_for_diff_stage.jpg "figure1")
+
+
+##### Comparisons to SOTA
+We compared **Wan2.1** with leading open-source and closed-source models to evaluate the performace. Using our carefully designed set of 1,035 internal prompts, we tested across 14 major dimensions and 26 sub-dimensions. Then we calculated the total score through a weighted average based on the importance of each dimension. The detailed results are shown in the table below. These results demonstrate our model's superior performance compared to both open-source and closed-source models.
+
+![figure1](assets/vben_vs_sota.png "figure1")
+
+
+## Citation
+If you find our work helpful, please cite us.
+
+```
+@article{wan2.1,
+    title   = {Wan: Open and Advanced Large-Scale Video Generative Models},
+    author  = {Wan Team},
+    journal = {},
+    year    = {2025}
+}
+```
+
+## License Agreement
+The models in this repository are licensed under the Apache 2.0 License. We claim no rights over the your generate contents, granting you the freedom to use them while ensuring that your usage complies with the provisions of this license. You are fully accountable for your use of the models, which must not involve sharing any content that violates applicable laws, causes harm to individuals or groups, disseminates personal information intended for harm, spreads misinformation, or targets vulnerable populations. For a complete list of restrictions and details regarding your rights, please refer to the full text of the [license](LICENSE.txt).
+
+
+## Acknowledgements
+
+We would like to thank the contributors to the [SD3](https://huggingface.co/stabilityai/stable-diffusion-3-medium), [Qwen](https://huggingface.co/Qwen), [umt5-xxl](https://huggingface.co/google/umt5-xxl), [diffusers](https://github.com/huggingface/diffusers) and [HuggingFace](https://huggingface.co) repositories, for their open research.
+
+
+
+## Contact Us
+If you would like to leave a message to our research or product teams, feel free to join our [Discord](https://discord.gg/p5XbdQV7) or [WeChat groups](https://gw.alicdn.com/imgextra/i2/O1CN01tqjWFi1ByuyehkTSB_!!6000000000015-0-tps-611-1279.jpg)!

wan_models/Wan2.1-T2V-1.3B/config.json

@@ -0,0 +1,14 @@
+{
+  "_class_name": "WanModel",
+  "_diffusers_version": "0.30.0",
+  "dim": 1536,
+  "eps": 1e-06,
+  "ffn_dim": 8960,
+  "freq_dim": 256,
+  "in_dim": 16,
+  "model_type": "t2v",
+  "num_heads": 12,
+  "num_layers": 30,
+  "out_dim": 16,
+  "text_len": 512
+}

wan_models/Wan2.1-T2V-1.3B/google/umt5-xxl/special_tokens_map.json

@@ -0,0 +1,308 @@
+{
+  "additional_special_tokens": [
+    "<extra_id_0>",
+    "<extra_id_1>",
+    "<extra_id_2>",
+    "<extra_id_3>",
+    "<extra_id_4>",
+    "<extra_id_5>",
+    "<extra_id_6>",
+    "<extra_id_7>",
+    "<extra_id_8>",
+    "<extra_id_9>",
+    "<extra_id_10>",
+    "<extra_id_11>",
+    "<extra_id_12>",
+    "<extra_id_13>",
+    "<extra_id_14>",
+    "<extra_id_15>",
+    "<extra_id_16>",
+    "<extra_id_17>",
+    "<extra_id_18>",
+    "<extra_id_19>",
+    "<extra_id_20>",
+    "<extra_id_21>",
+    "<extra_id_22>",
+    "<extra_id_23>",
+    "<extra_id_24>",
+    "<extra_id_25>",
+    "<extra_id_26>",
+    "<extra_id_27>",
+    "<extra_id_28>",
+    "<extra_id_29>",
+    "<extra_id_30>",
+    "<extra_id_31>",
+    "<extra_id_32>",
+    "<extra_id_33>",
+    "<extra_id_34>",
+    "<extra_id_35>",
+    "<extra_id_36>",
+    "<extra_id_37>",
+    "<extra_id_38>",
+    "<extra_id_39>",
+    "<extra_id_40>",
+    "<extra_id_41>",
+    "<extra_id_42>",
+    "<extra_id_43>",
+    "<extra_id_44>",
+    "<extra_id_45>",
+    "<extra_id_46>",
+    "<extra_id_47>",
+    "<extra_id_48>",
+    "<extra_id_49>",
+    "<extra_id_50>",
+    "<extra_id_51>",
+    "<extra_id_52>",
+    "<extra_id_53>",
+    "<extra_id_54>",
+    "<extra_id_55>",
+    "<extra_id_56>",
+    "<extra_id_57>",
+    "<extra_id_58>",
+    "<extra_id_59>",
+    "<extra_id_60>",
+    "<extra_id_61>",
+    "<extra_id_62>",
+    "<extra_id_63>",
+    "<extra_id_64>",
+    "<extra_id_65>",
+    "<extra_id_66>",
+    "<extra_id_67>",
+    "<extra_id_68>",
+    "<extra_id_69>",
+    "<extra_id_70>",
+    "<extra_id_71>",
+    "<extra_id_72>",
+    "<extra_id_73>",
+    "<extra_id_74>",
+    "<extra_id_75>",
+    "<extra_id_76>",
+    "<extra_id_77>",
+    "<extra_id_78>",
+    "<extra_id_79>",
+    "<extra_id_80>",
+    "<extra_id_81>",
+    "<extra_id_82>",
+    "<extra_id_83>",
+    "<extra_id_84>",
+    "<extra_id_85>",
+    "<extra_id_86>",
+    "<extra_id_87>",
+    "<extra_id_88>",
+    "<extra_id_89>",
+    "<extra_id_90>",
+    "<extra_id_91>",
+    "<extra_id_92>",
+    "<extra_id_93>",
+    "<extra_id_94>",
+    "<extra_id_95>",
+    "<extra_id_96>",
+    "<extra_id_97>",
+    "<extra_id_98>",
+    "<extra_id_99>",
+    "<extra_id_100>",
+    "<extra_id_101>",
+    "<extra_id_102>",
+    "<extra_id_103>",
+    "<extra_id_104>",
+    "<extra_id_105>",
+    "<extra_id_106>",
+    "<extra_id_107>",
+    "<extra_id_108>",
+    "<extra_id_109>",
+    "<extra_id_110>",
+    "<extra_id_111>",
+    "<extra_id_112>",
+    "<extra_id_113>",
+    "<extra_id_114>",
+    "<extra_id_115>",
+    "<extra_id_116>",
+    "<extra_id_117>",
+    "<extra_id_118>",
+    "<extra_id_119>",
+    "<extra_id_120>",
+    "<extra_id_121>",
+    "<extra_id_122>",
+    "<extra_id_123>",
+    "<extra_id_124>",
+    "<extra_id_125>",
+    "<extra_id_126>",
+    "<extra_id_127>",
+    "<extra_id_128>",
+    "<extra_id_129>",
+    "<extra_id_130>",
+    "<extra_id_131>",
+    "<extra_id_132>",
+    "<extra_id_133>",
+    "<extra_id_134>",
+    "<extra_id_135>",
+    "<extra_id_136>",
+    "<extra_id_137>",
+    "<extra_id_138>",
+    "<extra_id_139>",
+    "<extra_id_140>",
+    "<extra_id_141>",
+    "<extra_id_142>",
+    "<extra_id_143>",
+    "<extra_id_144>",
+    "<extra_id_145>",
+    "<extra_id_146>",
+    "<extra_id_147>",
+    "<extra_id_148>",
+    "<extra_id_149>",
+    "<extra_id_150>",
+    "<extra_id_151>",
+    "<extra_id_152>",
+    "<extra_id_153>",
+    "<extra_id_154>",
+    "<extra_id_155>",
+    "<extra_id_156>",
+    "<extra_id_157>",
+    "<extra_id_158>",
+    "<extra_id_159>",
+    "<extra_id_160>",
+    "<extra_id_161>",
+    "<extra_id_162>",
+    "<extra_id_163>",
+    "<extra_id_164>",
+    "<extra_id_165>",
+    "<extra_id_166>",
+    "<extra_id_167>",
+    "<extra_id_168>",
+    "<extra_id_169>",
+    "<extra_id_170>",
+    "<extra_id_171>",
+    "<extra_id_172>",
+    "<extra_id_173>",
+    "<extra_id_174>",
+    "<extra_id_175>",
+    "<extra_id_176>",
+    "<extra_id_177>",
+    "<extra_id_178>",
+    "<extra_id_179>",
+    "<extra_id_180>",
+    "<extra_id_181>",
+    "<extra_id_182>",
+    "<extra_id_183>",
+    "<extra_id_184>",
+    "<extra_id_185>",
+    "<extra_id_186>",
+    "<extra_id_187>",
+    "<extra_id_188>",
+    "<extra_id_189>",
+    "<extra_id_190>",
+    "<extra_id_191>",
+    "<extra_id_192>",
+    "<extra_id_193>",
+    "<extra_id_194>",
+    "<extra_id_195>",
+    "<extra_id_196>",
+    "<extra_id_197>",
+    "<extra_id_198>",
+    "<extra_id_199>",
+    "<extra_id_200>",
+    "<extra_id_201>",
+    "<extra_id_202>",
+    "<extra_id_203>",
+    "<extra_id_204>",
+    "<extra_id_205>",
+    "<extra_id_206>",
+    "<extra_id_207>",
+    "<extra_id_208>",
+    "<extra_id_209>",
+    "<extra_id_210>",
+    "<extra_id_211>",
+    "<extra_id_212>",
+    "<extra_id_213>",
+    "<extra_id_214>",
+    "<extra_id_215>",
+    "<extra_id_216>",
+    "<extra_id_217>",
+    "<extra_id_218>",
+    "<extra_id_219>",
+    "<extra_id_220>",
+    "<extra_id_221>",
+    "<extra_id_222>",
+    "<extra_id_223>",
+    "<extra_id_224>",
+    "<extra_id_225>",
+    "<extra_id_226>",
+    "<extra_id_227>",
+    "<extra_id_228>",
+    "<extra_id_229>",
+    "<extra_id_230>",
+    "<extra_id_231>",
+    "<extra_id_232>",
+    "<extra_id_233>",
+    "<extra_id_234>",
+    "<extra_id_235>",
+    "<extra_id_236>",
+    "<extra_id_237>",
+    "<extra_id_238>",
+    "<extra_id_239>",
+    "<extra_id_240>",
+    "<extra_id_241>",
+    "<extra_id_242>",
+    "<extra_id_243>",
+    "<extra_id_244>",
+    "<extra_id_245>",
+    "<extra_id_246>",
+    "<extra_id_247>",
+    "<extra_id_248>",
+    "<extra_id_249>",
+    "<extra_id_250>",
+    "<extra_id_251>",
+    "<extra_id_252>",
+    "<extra_id_253>",
+    "<extra_id_254>",
+    "<extra_id_255>",
+    "<extra_id_256>",
+    "<extra_id_257>",
+    "<extra_id_258>",
+    "<extra_id_259>",
+    "<extra_id_260>",
+    "<extra_id_261>",
+    "<extra_id_262>",
+    "<extra_id_263>",
+    "<extra_id_264>",
+    "<extra_id_265>",
+    "<extra_id_266>",
+    "<extra_id_267>",
+    "<extra_id_268>",
+    "<extra_id_269>",
+    "<extra_id_270>",
+    "<extra_id_271>",
+    "<extra_id_272>",
+    "<extra_id_273>",
+    "<extra_id_274>",
+    "<extra_id_275>",
+    "<extra_id_276>",
+    "<extra_id_277>",
+    "<extra_id_278>",
+    "<extra_id_279>",
+    "<extra_id_280>",
+    "<extra_id_281>",
+    "<extra_id_282>",
+    "<extra_id_283>",
+    "<extra_id_284>",
+    "<extra_id_285>",
+    "<extra_id_286>",
+    "<extra_id_287>",
+    "<extra_id_288>",
+    "<extra_id_289>",
+    "<extra_id_290>",
+    "<extra_id_291>",
+    "<extra_id_292>",
+    "<extra_id_293>",
+    "<extra_id_294>",
+    "<extra_id_295>",
+    "<extra_id_296>",
+    "<extra_id_297>",
+    "<extra_id_298>",
+    "<extra_id_299>"
+  ],
+  "bos_token": "<s>",
+  "eos_token": "</s>",
+  "pad_token": "<pad>",
+  "unk_token": "<unk>"
+}

wan_models/Wan2.1-T2V-1.3B/google/umt5-xxl/tokenizer_config.json

@@ -0,0 +1,2748 @@
+{
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<pad>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "</s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "<s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "3": {
+      "content": "<unk>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256000": {
+      "content": "<extra_id_299>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256001": {
+      "content": "<extra_id_298>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256002": {
+      "content": "<extra_id_297>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256003": {
+      "content": "<extra_id_296>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256004": {
+      "content": "<extra_id_295>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256005": {
+      "content": "<extra_id_294>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256006": {
+      "content": "<extra_id_293>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256007": {
+      "content": "<extra_id_292>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256008": {
+      "content": "<extra_id_291>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256009": {
+      "content": "<extra_id_290>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256010": {
+      "content": "<extra_id_289>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256011": {
+      "content": "<extra_id_288>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256012": {
+      "content": "<extra_id_287>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256013": {
+      "content": "<extra_id_286>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256014": {
+      "content": "<extra_id_285>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256015": {
+      "content": "<extra_id_284>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256016": {
+      "content": "<extra_id_283>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256017": {
+      "content": "<extra_id_282>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256018": {
+      "content": "<extra_id_281>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256019": {
+      "content": "<extra_id_280>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256020": {
+      "content": "<extra_id_279>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256021": {
+      "content": "<extra_id_278>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256022": {
+      "content": "<extra_id_277>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256023": {
+      "content": "<extra_id_276>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256024": {
+      "content": "<extra_id_275>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256025": {
+      "content": "<extra_id_274>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256026": {
+      "content": "<extra_id_273>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256027": {
+      "content": "<extra_id_272>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256028": {
+      "content": "<extra_id_271>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256029": {
+      "content": "<extra_id_270>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256030": {
+      "content": "<extra_id_269>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256031": {
+      "content": "<extra_id_268>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256032": {
+      "content": "<extra_id_267>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256033": {
+      "content": "<extra_id_266>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256034": {
+      "content": "<extra_id_265>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256035": {
+      "content": "<extra_id_264>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256036": {
+      "content": "<extra_id_263>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256037": {
+      "content": "<extra_id_262>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256038": {
+      "content": "<extra_id_261>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256039": {
+      "content": "<extra_id_260>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256040": {
+      "content": "<extra_id_259>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256041": {
+      "content": "<extra_id_258>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256042": {
+      "content": "<extra_id_257>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256043": {
+      "content": "<extra_id_256>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256044": {
+      "content": "<extra_id_255>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256045": {
+      "content": "<extra_id_254>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256046": {
+      "content": "<extra_id_253>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256047": {
+      "content": "<extra_id_252>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256048": {
+      "content": "<extra_id_251>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256049": {
+      "content": "<extra_id_250>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256050": {
+      "content": "<extra_id_249>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256051": {
+      "content": "<extra_id_248>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256052": {
+      "content": "<extra_id_247>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256053": {
+      "content": "<extra_id_246>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256054": {
+      "content": "<extra_id_245>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256055": {
+      "content": "<extra_id_244>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256056": {
+      "content": "<extra_id_243>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256057": {
+      "content": "<extra_id_242>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256058": {
+      "content": "<extra_id_241>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256059": {
+      "content": "<extra_id_240>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256060": {
+      "content": "<extra_id_239>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256061": {
+      "content": "<extra_id_238>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256062": {
+      "content": "<extra_id_237>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256063": {
+      "content": "<extra_id_236>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256064": {
+      "content": "<extra_id_235>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256065": {
+      "content": "<extra_id_234>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256066": {
+      "content": "<extra_id_233>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256067": {
+      "content": "<extra_id_232>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256068": {
+      "content": "<extra_id_231>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256069": {
+      "content": "<extra_id_230>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256070": {
+      "content": "<extra_id_229>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256071": {
+      "content": "<extra_id_228>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256072": {
+      "content": "<extra_id_227>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256073": {
+      "content": "<extra_id_226>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256074": {
+      "content": "<extra_id_225>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256075": {
+      "content": "<extra_id_224>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256076": {
+      "content": "<extra_id_223>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256077": {
+      "content": "<extra_id_222>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256078": {
+      "content": "<extra_id_221>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256079": {
+      "content": "<extra_id_220>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256080": {
+      "content": "<extra_id_219>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256081": {
+      "content": "<extra_id_218>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256082": {
+      "content": "<extra_id_217>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256083": {
+      "content": "<extra_id_216>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256084": {
+      "content": "<extra_id_215>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256085": {
+      "content": "<extra_id_214>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256086": {
+      "content": "<extra_id_213>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256087": {
+      "content": "<extra_id_212>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256088": {
+      "content": "<extra_id_211>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256089": {
+      "content": "<extra_id_210>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256090": {
+      "content": "<extra_id_209>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256091": {
+      "content": "<extra_id_208>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256092": {
+      "content": "<extra_id_207>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256093": {
+      "content": "<extra_id_206>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256094": {
+      "content": "<extra_id_205>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256095": {
+      "content": "<extra_id_204>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256096": {
+      "content": "<extra_id_203>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256097": {
+      "content": "<extra_id_202>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256098": {
+      "content": "<extra_id_201>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256099": {
+      "content": "<extra_id_200>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256100": {
+      "content": "<extra_id_199>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256101": {
+      "content": "<extra_id_198>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256102": {
+      "content": "<extra_id_197>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256103": {
+      "content": "<extra_id_196>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256104": {
+      "content": "<extra_id_195>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256105": {
+      "content": "<extra_id_194>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256106": {
+      "content": "<extra_id_193>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256107": {
+      "content": "<extra_id_192>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256108": {
+      "content": "<extra_id_191>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256109": {
+      "content": "<extra_id_190>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256110": {
+      "content": "<extra_id_189>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256111": {
+      "content": "<extra_id_188>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256112": {
+      "content": "<extra_id_187>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256113": {
+      "content": "<extra_id_186>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256114": {
+      "content": "<extra_id_185>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256115": {
+      "content": "<extra_id_184>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256116": {
+      "content": "<extra_id_183>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256117": {
+      "content": "<extra_id_182>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256118": {
+      "content": "<extra_id_181>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256119": {
+      "content": "<extra_id_180>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256120": {
+      "content": "<extra_id_179>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256121": {
+      "content": "<extra_id_178>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256122": {
+      "content": "<extra_id_177>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256123": {
+      "content": "<extra_id_176>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256124": {
+      "content": "<extra_id_175>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256125": {
+      "content": "<extra_id_174>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256126": {
+      "content": "<extra_id_173>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256127": {
+      "content": "<extra_id_172>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256128": {
+      "content": "<extra_id_171>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256129": {
+      "content": "<extra_id_170>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256130": {
+      "content": "<extra_id_169>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256131": {
+      "content": "<extra_id_168>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256132": {
+      "content": "<extra_id_167>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256133": {
+      "content": "<extra_id_166>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256134": {
+      "content": "<extra_id_165>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256135": {
+      "content": "<extra_id_164>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256136": {
+      "content": "<extra_id_163>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256137": {
+      "content": "<extra_id_162>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256138": {
+      "content": "<extra_id_161>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256139": {
+      "content": "<extra_id_160>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256140": {
+      "content": "<extra_id_159>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256141": {
+      "content": "<extra_id_158>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256142": {
+      "content": "<extra_id_157>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256143": {
+      "content": "<extra_id_156>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256144": {
+      "content": "<extra_id_155>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256145": {
+      "content": "<extra_id_154>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256146": {
+      "content": "<extra_id_153>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256147": {
+      "content": "<extra_id_152>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256148": {
+      "content": "<extra_id_151>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256149": {
+      "content": "<extra_id_150>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256150": {
+      "content": "<extra_id_149>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256151": {
+      "content": "<extra_id_148>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256152": {
+      "content": "<extra_id_147>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256153": {
+      "content": "<extra_id_146>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256154": {
+      "content": "<extra_id_145>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256155": {
+      "content": "<extra_id_144>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256156": {
+      "content": "<extra_id_143>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256157": {
+      "content": "<extra_id_142>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256158": {
+      "content": "<extra_id_141>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256159": {
+      "content": "<extra_id_140>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256160": {
+      "content": "<extra_id_139>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256161": {
+      "content": "<extra_id_138>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256162": {
+      "content": "<extra_id_137>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256163": {
+      "content": "<extra_id_136>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256164": {
+      "content": "<extra_id_135>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256165": {
+      "content": "<extra_id_134>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256166": {
+      "content": "<extra_id_133>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256167": {
+      "content": "<extra_id_132>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256168": {
+      "content": "<extra_id_131>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256169": {
+      "content": "<extra_id_130>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256170": {
+      "content": "<extra_id_129>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256171": {
+      "content": "<extra_id_128>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256172": {
+      "content": "<extra_id_127>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256173": {
+      "content": "<extra_id_126>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256174": {
+      "content": "<extra_id_125>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256175": {
+      "content": "<extra_id_124>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256176": {
+      "content": "<extra_id_123>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256177": {
+      "content": "<extra_id_122>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256178": {
+      "content": "<extra_id_121>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256179": {
+      "content": "<extra_id_120>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256180": {
+      "content": "<extra_id_119>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256181": {
+      "content": "<extra_id_118>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256182": {
+      "content": "<extra_id_117>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256183": {
+      "content": "<extra_id_116>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256184": {
+      "content": "<extra_id_115>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256185": {
+      "content": "<extra_id_114>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256186": {
+      "content": "<extra_id_113>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256187": {
+      "content": "<extra_id_112>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256188": {
+      "content": "<extra_id_111>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256189": {
+      "content": "<extra_id_110>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256190": {
+      "content": "<extra_id_109>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256191": {
+      "content": "<extra_id_108>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256192": {
+      "content": "<extra_id_107>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256193": {
+      "content": "<extra_id_106>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256194": {
+      "content": "<extra_id_105>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256195": {
+      "content": "<extra_id_104>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256196": {
+      "content": "<extra_id_103>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256197": {
+      "content": "<extra_id_102>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256198": {
+      "content": "<extra_id_101>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256199": {
+      "content": "<extra_id_100>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256200": {
+      "content": "<extra_id_99>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256201": {
+      "content": "<extra_id_98>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256202": {
+      "content": "<extra_id_97>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256203": {
+      "content": "<extra_id_96>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256204": {
+      "content": "<extra_id_95>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256205": {
+      "content": "<extra_id_94>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256206": {
+      "content": "<extra_id_93>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256207": {
+      "content": "<extra_id_92>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256208": {
+      "content": "<extra_id_91>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256209": {
+      "content": "<extra_id_90>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256210": {
+      "content": "<extra_id_89>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256211": {
+      "content": "<extra_id_88>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256212": {
+      "content": "<extra_id_87>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256213": {
+      "content": "<extra_id_86>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256214": {
+      "content": "<extra_id_85>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256215": {
+      "content": "<extra_id_84>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256216": {
+      "content": "<extra_id_83>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256217": {
+      "content": "<extra_id_82>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256218": {
+      "content": "<extra_id_81>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256219": {
+      "content": "<extra_id_80>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256220": {
+      "content": "<extra_id_79>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256221": {
+      "content": "<extra_id_78>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256222": {
+      "content": "<extra_id_77>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256223": {
+      "content": "<extra_id_76>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256224": {
+      "content": "<extra_id_75>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256225": {
+      "content": "<extra_id_74>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256226": {
+      "content": "<extra_id_73>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256227": {
+      "content": "<extra_id_72>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256228": {
+      "content": "<extra_id_71>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256229": {
+      "content": "<extra_id_70>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256230": {
+      "content": "<extra_id_69>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256231": {
+      "content": "<extra_id_68>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256232": {
+      "content": "<extra_id_67>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256233": {
+      "content": "<extra_id_66>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256234": {
+      "content": "<extra_id_65>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256235": {
+      "content": "<extra_id_64>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256236": {
+      "content": "<extra_id_63>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256237": {
+      "content": "<extra_id_62>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256238": {
+      "content": "<extra_id_61>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256239": {
+      "content": "<extra_id_60>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256240": {
+      "content": "<extra_id_59>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256241": {
+      "content": "<extra_id_58>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256242": {
+      "content": "<extra_id_57>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256243": {
+      "content": "<extra_id_56>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256244": {
+      "content": "<extra_id_55>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256245": {
+      "content": "<extra_id_54>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256246": {
+      "content": "<extra_id_53>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256247": {
+      "content": "<extra_id_52>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256248": {
+      "content": "<extra_id_51>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256249": {
+      "content": "<extra_id_50>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256250": {
+      "content": "<extra_id_49>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256251": {
+      "content": "<extra_id_48>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256252": {
+      "content": "<extra_id_47>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256253": {
+      "content": "<extra_id_46>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256254": {
+      "content": "<extra_id_45>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256255": {
+      "content": "<extra_id_44>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256256": {
+      "content": "<extra_id_43>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256257": {
+      "content": "<extra_id_42>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256258": {
+      "content": "<extra_id_41>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256259": {
+      "content": "<extra_id_40>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256260": {
+      "content": "<extra_id_39>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256261": {
+      "content": "<extra_id_38>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256262": {
+      "content": "<extra_id_37>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256263": {
+      "content": "<extra_id_36>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256264": {
+      "content": "<extra_id_35>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256265": {
+      "content": "<extra_id_34>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256266": {
+      "content": "<extra_id_33>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256267": {
+      "content": "<extra_id_32>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256268": {
+      "content": "<extra_id_31>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256269": {
+      "content": "<extra_id_30>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256270": {
+      "content": "<extra_id_29>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256271": {
+      "content": "<extra_id_28>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256272": {
+      "content": "<extra_id_27>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256273": {
+      "content": "<extra_id_26>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256274": {
+      "content": "<extra_id_25>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256275": {
+      "content": "<extra_id_24>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256276": {
+      "content": "<extra_id_23>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256277": {
+      "content": "<extra_id_22>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256278": {
+      "content": "<extra_id_21>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256279": {
+      "content": "<extra_id_20>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256280": {
+      "content": "<extra_id_19>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256281": {
+      "content": "<extra_id_18>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256282": {
+      "content": "<extra_id_17>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256283": {
+      "content": "<extra_id_16>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256284": {
+      "content": "<extra_id_15>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256285": {
+      "content": "<extra_id_14>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256286": {
+      "content": "<extra_id_13>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256287": {
+      "content": "<extra_id_12>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256288": {
+      "content": "<extra_id_11>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256289": {
+      "content": "<extra_id_10>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256290": {
+      "content": "<extra_id_9>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256291": {
+      "content": "<extra_id_8>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256292": {
+      "content": "<extra_id_7>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256293": {
+      "content": "<extra_id_6>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256294": {
+      "content": "<extra_id_5>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256295": {
+      "content": "<extra_id_4>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256296": {
+      "content": "<extra_id_3>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256297": {
+      "content": "<extra_id_2>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256298": {
+      "content": "<extra_id_1>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256299": {
+      "content": "<extra_id_0>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "additional_special_tokens": [
+    "<extra_id_0>",
+    "<extra_id_1>",
+    "<extra_id_2>",
+    "<extra_id_3>",
+    "<extra_id_4>",
+    "<extra_id_5>",
+    "<extra_id_6>",
+    "<extra_id_7>",
+    "<extra_id_8>",
+    "<extra_id_9>",
+    "<extra_id_10>",
+    "<extra_id_11>",
+    "<extra_id_12>",
+    "<extra_id_13>",
+    "<extra_id_14>",
+    "<extra_id_15>",
+    "<extra_id_16>",
+    "<extra_id_17>",
+    "<extra_id_18>",
+    "<extra_id_19>",
+    "<extra_id_20>",
+    "<extra_id_21>",
+    "<extra_id_22>",
+    "<extra_id_23>",
+    "<extra_id_24>",
+    "<extra_id_25>",
+    "<extra_id_26>",
+    "<extra_id_27>",
+    "<extra_id_28>",
+    "<extra_id_29>",
+    "<extra_id_30>",
+    "<extra_id_31>",
+    "<extra_id_32>",
+    "<extra_id_33>",
+    "<extra_id_34>",
+    "<extra_id_35>",
+    "<extra_id_36>",
+    "<extra_id_37>",
+    "<extra_id_38>",
+    "<extra_id_39>",
+    "<extra_id_40>",
+    "<extra_id_41>",
+    "<extra_id_42>",
+    "<extra_id_43>",
+    "<extra_id_44>",
+    "<extra_id_45>",
+    "<extra_id_46>",
+    "<extra_id_47>",
+    "<extra_id_48>",
+    "<extra_id_49>",
+    "<extra_id_50>",
+    "<extra_id_51>",
+    "<extra_id_52>",
+    "<extra_id_53>",
+    "<extra_id_54>",
+    "<extra_id_55>",
+    "<extra_id_56>",
+    "<extra_id_57>",
+    "<extra_id_58>",
+    "<extra_id_59>",
+    "<extra_id_60>",
+    "<extra_id_61>",
+    "<extra_id_62>",
+    "<extra_id_63>",
+    "<extra_id_64>",
+    "<extra_id_65>",
+    "<extra_id_66>",
+    "<extra_id_67>",
+    "<extra_id_68>",
+    "<extra_id_69>",
+    "<extra_id_70>",
+    "<extra_id_71>",
+    "<extra_id_72>",
+    "<extra_id_73>",
+    "<extra_id_74>",
+    "<extra_id_75>",
+    "<extra_id_76>",
+    "<extra_id_77>",
+    "<extra_id_78>",
+    "<extra_id_79>",
+    "<extra_id_80>",
+    "<extra_id_81>",
+    "<extra_id_82>",
+    "<extra_id_83>",
+    "<extra_id_84>",
+    "<extra_id_85>",
+    "<extra_id_86>",
+    "<extra_id_87>",
+    "<extra_id_88>",
+    "<extra_id_89>",
+    "<extra_id_90>",
+    "<extra_id_91>",
+    "<extra_id_92>",
+    "<extra_id_93>",
+    "<extra_id_94>",
+    "<extra_id_95>",
+    "<extra_id_96>",
+    "<extra_id_97>",
+    "<extra_id_98>",
+    "<extra_id_99>",
+    "<extra_id_100>",
+    "<extra_id_101>",
+    "<extra_id_102>",
+    "<extra_id_103>",
+    "<extra_id_104>",
+    "<extra_id_105>",
+    "<extra_id_106>",
+    "<extra_id_107>",
+    "<extra_id_108>",
+    "<extra_id_109>",
+    "<extra_id_110>",
+    "<extra_id_111>",
+    "<extra_id_112>",
+    "<extra_id_113>",
+    "<extra_id_114>",
+    "<extra_id_115>",
+    "<extra_id_116>",
+    "<extra_id_117>",
+    "<extra_id_118>",
+    "<extra_id_119>",
+    "<extra_id_120>",
+    "<extra_id_121>",
+    "<extra_id_122>",
+    "<extra_id_123>",
+    "<extra_id_124>",
+    "<extra_id_125>",
+    "<extra_id_126>",
+    "<extra_id_127>",
+    "<extra_id_128>",
+    "<extra_id_129>",
+    "<extra_id_130>",
+    "<extra_id_131>",
+    "<extra_id_132>",
+    "<extra_id_133>",
+    "<extra_id_134>",
+    "<extra_id_135>",
+    "<extra_id_136>",
+    "<extra_id_137>",
+    "<extra_id_138>",
+    "<extra_id_139>",
+    "<extra_id_140>",
+    "<extra_id_141>",
+    "<extra_id_142>",
+    "<extra_id_143>",
+    "<extra_id_144>",
+    "<extra_id_145>",
+    "<extra_id_146>",
+    "<extra_id_147>",
+    "<extra_id_148>",
+    "<extra_id_149>",
+    "<extra_id_150>",
+    "<extra_id_151>",
+    "<extra_id_152>",
+    "<extra_id_153>",
+    "<extra_id_154>",
+    "<extra_id_155>",
+    "<extra_id_156>",
+    "<extra_id_157>",
+    "<extra_id_158>",
+    "<extra_id_159>",
+    "<extra_id_160>",
+    "<extra_id_161>",
+    "<extra_id_162>",
+    "<extra_id_163>",
+    "<extra_id_164>",
+    "<extra_id_165>",
+    "<extra_id_166>",
+    "<extra_id_167>",
+    "<extra_id_168>",
+    "<extra_id_169>",
+    "<extra_id_170>",
+    "<extra_id_171>",
+    "<extra_id_172>",
+    "<extra_id_173>",
+    "<extra_id_174>",
+    "<extra_id_175>",
+    "<extra_id_176>",
+    "<extra_id_177>",
+    "<extra_id_178>",
+    "<extra_id_179>",
+    "<extra_id_180>",
+    "<extra_id_181>",
+    "<extra_id_182>",
+    "<extra_id_183>",
+    "<extra_id_184>",
+    "<extra_id_185>",
+    "<extra_id_186>",
+    "<extra_id_187>",
+    "<extra_id_188>",
+    "<extra_id_189>",
+    "<extra_id_190>",
+    "<extra_id_191>",
+    "<extra_id_192>",
+    "<extra_id_193>",
+    "<extra_id_194>",
+    "<extra_id_195>",
+    "<extra_id_196>",
+    "<extra_id_197>",
+    "<extra_id_198>",
+    "<extra_id_199>",
+    "<extra_id_200>",
+    "<extra_id_201>",
+    "<extra_id_202>",
+    "<extra_id_203>",
+    "<extra_id_204>",
+    "<extra_id_205>",
+    "<extra_id_206>",
+    "<extra_id_207>",
+    "<extra_id_208>",
+    "<extra_id_209>",
+    "<extra_id_210>",
+    "<extra_id_211>",
+    "<extra_id_212>",
+    "<extra_id_213>",
+    "<extra_id_214>",
+    "<extra_id_215>",
+    "<extra_id_216>",
+    "<extra_id_217>",
+    "<extra_id_218>",
+    "<extra_id_219>",
+    "<extra_id_220>",
+    "<extra_id_221>",
+    "<extra_id_222>",
+    "<extra_id_223>",
+    "<extra_id_224>",
+    "<extra_id_225>",
+    "<extra_id_226>",
+    "<extra_id_227>",
+    "<extra_id_228>",
+    "<extra_id_229>",
+    "<extra_id_230>",
+    "<extra_id_231>",
+    "<extra_id_232>",
+    "<extra_id_233>",
+    "<extra_id_234>",
+    "<extra_id_235>",
+    "<extra_id_236>",
+    "<extra_id_237>",
+    "<extra_id_238>",
+    "<extra_id_239>",
+    "<extra_id_240>",
+    "<extra_id_241>",
+    "<extra_id_242>",
+    "<extra_id_243>",
+    "<extra_id_244>",
+    "<extra_id_245>",
+    "<extra_id_246>",
+    "<extra_id_247>",
+    "<extra_id_248>",
+    "<extra_id_249>",
+    "<extra_id_250>",
+    "<extra_id_251>",
+    "<extra_id_252>",
+    "<extra_id_253>",
+    "<extra_id_254>",
+    "<extra_id_255>",
+    "<extra_id_256>",
+    "<extra_id_257>",
+    "<extra_id_258>",
+    "<extra_id_259>",
+    "<extra_id_260>",
+    "<extra_id_261>",
+    "<extra_id_262>",
+    "<extra_id_263>",
+    "<extra_id_264>",
+    "<extra_id_265>",
+    "<extra_id_266>",
+    "<extra_id_267>",
+    "<extra_id_268>",
+    "<extra_id_269>",
+    "<extra_id_270>",
+    "<extra_id_271>",
+    "<extra_id_272>",
+    "<extra_id_273>",
+    "<extra_id_274>",
+    "<extra_id_275>",
+    "<extra_id_276>",
+    "<extra_id_277>",
+    "<extra_id_278>",
+    "<extra_id_279>",
+    "<extra_id_280>",
+    "<extra_id_281>",
+    "<extra_id_282>",
+    "<extra_id_283>",
+    "<extra_id_284>",
+    "<extra_id_285>",
+    "<extra_id_286>",
+    "<extra_id_287>",
+    "<extra_id_288>",
+    "<extra_id_289>",
+    "<extra_id_290>",
+    "<extra_id_291>",
+    "<extra_id_292>",
+    "<extra_id_293>",
+    "<extra_id_294>",
+    "<extra_id_295>",
+    "<extra_id_296>",
+    "<extra_id_297>",
+    "<extra_id_298>",
+    "<extra_id_299>"
+  ],
+  "bos_token": "<s>",
+  "clean_up_tokenization_spaces": true,
+  "eos_token": "</s>",
+  "extra_ids": 300,
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "<pad>",
+  "sp_model_kwargs": {},
+  "spaces_between_special_tokens": false,
+  "tokenizer_class": "T5Tokenizer",
+  "unk_token": "<unk>"
+}

wan_models/Wan2.1-T2V-14B/LICENSE.txt

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   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.

wan_models/Wan2.1-T2V-14B/README.md

@@ -0,0 +1,301 @@
+---
+license: apache-2.0
+language:
+- en
+- zh
+pipeline_tag: text-to-video
+tags:
+- video generation
+library_name: diffusers
+inference:
+  parameters:
+    num_inference_steps: 10
+---
+# Wan2.1
+
+<p align="center">
+    <img src="assets/logo.png" width="400"/>
+<p>
+
+<p align="center">
+    💜 <a href=""><b>Wan</b></a> &nbsp&nbsp ｜ &nbsp&nbsp 🖥️ <a href="https://github.com/Wan-Video/Wan2.1">GitHub</a> &nbsp&nbsp  | &nbsp&nbsp🤗 <a href="https://huggingface.co/Wan-AI/">Hugging Face</a>&nbsp&nbsp | &nbsp&nbsp🤖 <a href="https://modelscope.cn/organization/Wan-AI">ModelScope</a>&nbsp&nbsp | &nbsp&nbsp 📑 <a href="">Paper (Coming soon)</a> &nbsp&nbsp | &nbsp&nbsp 📑 <a href="https://wanxai.com">Blog</a> &nbsp&nbsp | &nbsp&nbsp💬 <a href="https://gw.alicdn.com/imgextra/i2/O1CN01tqjWFi1ByuyehkTSB_!!6000000000015-0-tps-611-1279.jpg">WeChat Group</a>&nbsp&nbsp | &nbsp&nbsp 📖 <a href="https://discord.gg/p5XbdQV7">Discord</a>&nbsp&nbsp
+<br>
+
+-----
+
+[**Wan: Open and Advanced Large-Scale Video Generative Models**]("") <be>
+
+In this repository, we present **Wan2.1**, a comprehensive and open suite of video foundation models that pushes the boundaries of video generation. **Wan2.1** offers these key features:
+- 👍 **SOTA Performance**: **Wan2.1** consistently outperforms existing open-source models and state-of-the-art commercial solutions across multiple benchmarks.
+- 👍 **Supports Consumer-grade GPUs**: The T2V-1.3B model requires only 8.19 GB VRAM, making it compatible with almost all consumer-grade GPUs. It can generate a 5-second 480P video on an RTX 4090 in about 4 minutes (without optimization techniques like quantization). Its performance is even comparable to some closed-source models.
+- 👍 **Multiple Tasks**: **Wan2.1** excels in Text-to-Video, Image-to-Video, Video Editing, Text-to-Image, and Video-to-Audio, advancing the field of video generation.
+- 👍 **Visual Text Generation**: **Wan2.1** is the first video model capable of generating both Chinese and English text, featuring robust text generation that enhances its practical applications.
+- 👍 **Powerful Video VAE**: **Wan-VAE** delivers exceptional efficiency and performance, encoding and decoding 1080P videos of any length while preserving temporal information, making it an ideal foundation for video and image generation.
+
+This repository features our T2V-14B model, which establishes a new SOTA performance benchmark among both open-source and closed-source models. It demonstrates exceptional capabilities in generating high-quality visuals with significant motion dynamics. It is also the only video model capable of producing both Chinese and English text and supports video generation at both 480P and 720P resolutions.
+
+
+## Video Demos
+
+<div align="center">
+    <video width="80%" controls>
+        <source src="https://cloud.video.taobao.com/vod/Jth64Y7wNoPcJki_Bo1ZJTDBvNjsgjlVKsNs05Fqfps.mp4" type="video/mp4">
+        Your browser does not support the video tag.
+    </video>
+</div>
+
+## 🔥 Latest News!!
+
+* Feb 22, 2025: 👋 We've released the inference code and weights of Wan2.1.
+
+
+## 📑 Todo List
+- Wan2.1 Text-to-Video
+    - [x] Multi-GPU Inference code of the 14B and 1.3B models
+    - [x] Checkpoints of the 14B and 1.3B models
+    - [x] Gradio demo
+    - [ ] Diffusers integration
+    - [ ] ComfyUI integration
+- Wan2.1 Image-to-Video
+    - [x] Multi-GPU Inference code of the 14B model
+    - [x] Checkpoints of the 14B model
+    - [x] Gradio demo
+    - [ ] Diffusers integration
+    - [ ] ComfyUI integration
+
+
+## Quickstart
+
+#### Installation
+Clone the repo:
+```
+git clone https://github.com/Wan-Video/Wan2.1.git
+cd Wan2.1
+```
+
+Install dependencies:
+```
+# Ensure torch >= 2.4.0
+pip install -r requirements.txt
+```
+
+
+#### Model Download
+
+| Models        |                       Download Link                                           |    Notes                      |
+| --------------|-------------------------------------------------------------------------------|-------------------------------|
+| T2V-14B       |      🤗 [Huggingface](https://huggingface.co/Wan-AI/Wan2.1-T2V-14B)      🤖 [ModelScope](https://www.modelscope.cn/models/Wan-AI/Wan2.1-T2V-14B)          | Supports both 480P and 720P
+| I2V-14B-720P  |      🤗 [Huggingface](https://huggingface.co/Wan-AI/Wan2.1-I2V-14B-720P)    🤖 [ModelScope](https://www.modelscope.cn/models/Wan-AI/Wan2.1-I2V-14B-720P)     | Supports 720P
+| I2V-14B-480P  |      🤗 [Huggingface](https://huggingface.co/Wan-AI/Wan2.1-I2V-14B-480P)    🤖 [ModelScope](https://www.modelscope.cn/models/Wan-AI/Wan2.1-I2V-14B-480P)      | Supports 480P
+| T2V-1.3B      |      🤗 [Huggingface](https://huggingface.co/Wan-AI/Wan2.1-T2V-1.3B)     🤖 [ModelScope](https://www.modelscope.cn/models/Wan-AI/Wan2.1-T2V-1.3B)         | Supports 480P
+
+> 💡Note: The 1.3B model is capable of generating videos at 720P resolution. However, due to limited training at this resolution, the results are generally less stable compared to 480P. For optimal performance, we recommend using 480P resolution.
+
+
+Download models using 🤗 huggingface-cli:
+```
+pip install "huggingface_hub[cli]"
+huggingface-cli download Wan-AI/Wan2.1-T2V-14B --local-dir ./Wan2.1-T2V-14B
+```
+
+Download models using 🤖 modelscope-cli:
+```
+pip install modelscope
+modelscope download Wan-AI/Wan2.1-T2V-14B --local_dir ./Wan2.1-T2V-14B
+```
+#### Run Text-to-Video Generation
+
+This repository supports two Text-to-Video models (1.3B and 14B) and two resolutions (480P and 720P). The parameters and configurations for these models are as follows:
+
+<table>
+    <thead>
+        <tr>
+            <th rowspan="2">Task</th>
+            <th colspan="2">Resolution</th>
+            <th rowspan="2">Model</th>
+        </tr>
+        <tr>
+            <th>480P</th>
+            <th>720P</th>
+        </tr>
+    </thead>
+    <tbody>
+        <tr>
+            <td>t2v-14B</td>
+            <td style="color: green;">✔️</td>
+            <td style="color: green;">✔️</td>
+            <td>Wan2.1-T2V-14B</td>
+        </tr>
+        <tr>
+            <td>t2v-1.3B</td>
+            <td style="color: green;">✔️</td>
+            <td style="color: red;">❌</td>
+            <td>Wan2.1-T2V-1.3B</td>
+        </tr>
+    </tbody>
+</table>
+
+
+##### (1) Without Prompt Extention
+
+To facilitate implementation, we will start with a basic version of the inference process that skips the [prompt extension](#2-using-prompt-extention) step.
+
+- Single-GPU inference
+
+```
+python generate.py  --task t2v-14B --size 1280*720 --ckpt_dir ./Wan2.1-T2V-14B --prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage."
+```
+
+If you encounter OOM (Out-of-Memory) issues, you can use the `--offload_model True` and `--t5_cpu` options to reduce GPU memory usage. For example, on an RTX 4090 GPU:
+
+```
+python generate.py  --task t2v-1.3B --size 832*480 --ckpt_dir ./Wan2.1-T2V-1.3B --offload_model True --t5_cpu --sample_shift 8 --sample_guide_scale 6 --prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage."
+```
+
+> 💡Note: If you are using the `T2V-1.3B` model, we recommend setting the parameter `--sample_guide_scale 6`. The `--sample_shift parameter` can be adjusted within the range of 8 to 12 based on the performance.
+
+
+- Multi-GPU inference using FSDP + xDiT USP
+
+```
+pip install "xfuser>=0.4.1"
+torchrun --nproc_per_node=8 generate.py --task t2v-14B --size 1280*720 --ckpt_dir ./Wan2.1-T2V-14B --dit_fsdp --t5_fsdp --ulysses_size 8 --prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage."
+```
+
+
+##### (2) Using Prompt Extention
+
+Extending the prompts can effectively enrich the details in the generated videos, further enhancing the video quality. Therefore, we recommend enabling prompt extension. We provide the following two methods for prompt extension:
+
+- Use the Dashscope API for extension.
+  - Apply for a `dashscope.api_key` in advance ([EN](https://www.alibabacloud.com/help/en/model-studio/getting-started/first-api-call-to-qwen) | [CN](https://help.aliyun.com/zh/model-studio/getting-started/first-api-call-to-qwen)).
+  - Configure the environment variable `DASH_API_KEY` to specify the Dashscope API key. For users of Alibaba Cloud's international site, you also need to set the environment variable `DASH_API_URL` to 'https://dashscope-intl.aliyuncs.com/api/v1'. For more detailed instructions, please refer to the [dashscope document](https://www.alibabacloud.com/help/en/model-studio/developer-reference/use-qwen-by-calling-api?spm=a2c63.p38356.0.i1).
+  - Use the `qwen-plus` model for text-to-video tasks and `qwen-vl-max` for image-to-video tasks.
+  - You can modify the model used for extension with the parameter `--prompt_extend_model`. For example:
+```
+DASH_API_KEY=your_key python generate.py  --task t2v-14B --size 1280*720 --ckpt_dir ./Wan2.1-T2V-14B --prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage" --use_prompt_extend --prompt_extend_method 'dashscope' --prompt_extend_target_lang 'ch'
+```
+
+- Using a local model for extension.
+
+  - By default, the Qwen model on HuggingFace is used for this extension. Users can choose based on the available GPU memory size.
+  - For text-to-video tasks, you can use models like `Qwen/Qwen2.5-14B-Instruct`, `Qwen/Qwen2.5-7B-Instruct` and `Qwen/Qwen2.5-3B-Instruct`
+  - For image-to-video tasks, you can use models like `Qwen/Qwen2.5-VL-7B-Instruct` and `Qwen/Qwen2.5-VL-3B-Instruct`.
+  - Larger models generally provide better extension results but require more GPU memory.
+  - You can modify the model used for extension with the parameter `--prompt_extend_model` , allowing you to specify either a local model path or a Hugging Face model. For example:
+
+```
+python generate.py  --task t2v-14B --size 1280*720 --ckpt_dir ./Wan2.1-T2V-14B --prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage" --use_prompt_extend --prompt_extend_method 'local_qwen' --prompt_extend_target_lang 'ch'
+```
+
+##### (3) Runing local gradio
+
+```
+cd gradio
+# if one uses dashscope’s API for prompt extension
+DASH_API_KEY=your_key python t2v_14B_singleGPU.py --prompt_extend_method 'dashscope' --ckpt_dir ./Wan2.1-T2V-14B
+
+# if one uses a local model for prompt extension
+python t2v_14B_singleGPU.py --prompt_extend_method 'local_qwen' --ckpt_dir ./Wan2.1-T2V-14B
+```
+
+
+## Manual Evaluation
+
+
+Through manual evaluation, the results generated after prompt extension are superior to those from both closed-source and open-source models.
+
+<div align="center">
+    <img src="assets/t2v_res.jpg" alt="" style="width: 80%;" />
+</div>
+
+
+
+## Computational Efficiency on Different GPUs
+
+We test the computational efficiency of different **Wan2.1** models on different GPUs in the following table. The results are presented in the format: **Total time (s) / peak GPU memory (GB)**.
+
+
+<div align="center">
+    <img src="assets/comp_effic.png" alt="" style="width: 80%;" />
+</div>
+
+> The parameter settings for the tests presented in this table are as follows:
+> (1) For the 1.3B model on 8 GPUs, set `--ring_size 8` and `--ulysses_size 1`;
+> (2) For the 14B model on 1 GPU, use `--offload_model True`;
+> (3) For the 1.3B model on a single 4090 GPU, set `--offload_model True --t5_cpu`;
+> (4) For all testings, no prompt extension was applied, meaning `--use_prompt_extend` was not enabled.
+
+
+## Community Contributions
+- [DiffSynth-Studio](https://github.com/modelscope/DiffSynth-Studio) provides more support for Wan, including video-to-video, FP8 quantization, VRAM optimization, LoRA training, and more. Please refer to [their examples](https://github.com/modelscope/DiffSynth-Studio/tree/main/examples/wanvideo).
+
+-------
+
+## Introduction of Wan2.1
+
+**Wan2.1**  is designed on the mainstream diffusion transformer paradigm, achieving significant advancements in generative capabilities through a series of innovations. These include our novel spatio-temporal variational autoencoder (VAE), scalable training strategies, large-scale data construction, and automated evaluation metrics. Collectively, these contributions enhance the model’s performance and versatility.
+
+
+##### (1) 3D Variational Autoencoders
+We propose a novel 3D causal VAE architecture, termed **Wan-VAE** specifically designed for video generation. By combining multiple strategies, we improve spatio-temporal compression, reduce memory usage, and ensure temporal causality. **Wan-VAE** demonstrates significant advantages in performance efficiency compared to other open-source VAEs. Furthermore, our **Wan-VAE** can encode and decode unlimited-length 1080P videos without losing historical temporal information, making it particularly well-suited for video generation tasks.
+
+
+<div align="center">
+    <img src="assets/video_vae_res.jpg" alt="" style="width: 80%;" />
+</div>
+
+
+##### (2) Video Diffusion DiT
+
+**Wan2.1** is designed using the Flow Matching framework within the paradigm of mainstream Diffusion Transformers. Our model's architecture uses the T5 Encoder to encode multilingual text input, with cross-attention in each transformer block embedding the text into the model structure. Additionally, we employ an MLP with a Linear layer and a SiLU layer to process the input time embeddings and predict six modulation parameters individually. This MLP is shared across all transformer blocks, with each block learning a distinct set of biases. Our experimental findings reveal a significant performance improvement with this approach at the same parameter scale.
+
+<div align="center">
+    <img src="assets/video_dit_arch.jpg" alt="" style="width: 80%;" />
+</div>
+
+
+| Model  | Dimension | Input Dimension | Output Dimension | Feedforward Dimension | Frequency Dimension | Number of Heads | Number of Layers |
+|--------|-----------|-----------------|------------------|-----------------------|---------------------|-----------------|------------------|
+| 1.3B   | 1536      | 16              | 16               | 8960                  | 256                 | 12              | 30               |
+| 14B   | 5120       | 16              | 16               | 13824                 | 256                 | 40              | 40               |
+
+
+
+##### Data
+
+We curated and deduplicated a candidate dataset comprising a vast amount of image and video data. During the data curation process, we designed a four-step data cleaning process, focusing on fundamental dimensions, visual quality and motion quality. Through the robust data processing pipeline, we can easily obtain high-quality, diverse, and large-scale training sets of images and videos.
+
+![figure1](assets/data_for_diff_stage.jpg "figure1")
+
+
+##### Comparisons to SOTA
+We compared **Wan2.1** with leading open-source and closed-source models to evaluate the performace. Using our carefully designed set of 1,035 internal prompts, we tested across 14 major dimensions and 26 sub-dimensions. We then compute the total score by performing a weighted calculation on the scores of each dimension, utilizing weights derived from human preferences in the matching process. The detailed results are shown in the table below. These results demonstrate our model's superior performance compared to both open-source and closed-source models.
+
+![figure1](assets/vben_vs_sota.png "figure1")
+
+
+## Citation
+If you find our work helpful, please cite us.
+
+```
+@article{wan2.1,
+    title   = {Wan: Open and Advanced Large-Scale Video Generative Models},
+    author  = {Wan Team},
+    journal = {},
+    year    = {2025}
+}
+```
+
+## License Agreement
+The models in this repository are licensed under the Apache 2.0 License. We claim no rights over the your generate contents, granting you the freedom to use them while ensuring that your usage complies with the provisions of this license. You are fully accountable for your use of the models, which must not involve sharing any content that violates applicable laws, causes harm to individuals or groups, disseminates personal information intended for harm, spreads misinformation, or targets vulnerable populations. For a complete list of restrictions and details regarding your rights, please refer to the full text of the [license](LICENSE.txt).
+
+
+## Acknowledgements
+
+We would like to thank the contributors to the [SD3](https://huggingface.co/stabilityai/stable-diffusion-3-medium), [Qwen](https://huggingface.co/Qwen), [umt5-xxl](https://huggingface.co/google/umt5-xxl), [diffusers](https://github.com/huggingface/diffusers) and [HuggingFace](https://huggingface.co) repositories, for their open research.
+
+
+
+## Contact Us
+If you would like to leave a message to our research or product teams, feel free to join our [Discord](https://discord.gg/p5XbdQV7) or [WeChat groups](https://gw.alicdn.com/imgextra/i2/O1CN01tqjWFi1ByuyehkTSB_!!6000000000015-0-tps-611-1279.jpg)!

wan_models/Wan2.1-T2V-14B/diffusion_pytorch_model.safetensors.index.json

@@ -0,0 +1,1102 @@
+{
+  "metadata": {
+    "total_size": 57153966336
+  },
+  "weight_map": {
+    "blocks.0.cross_attn.k.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.0.cross_attn.k.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.0.cross_attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.0.cross_attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.0.cross_attn.o.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.0.cross_attn.o.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.0.cross_attn.q.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.0.cross_attn.q.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.0.cross_attn.v.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.0.cross_attn.v.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.0.ffn.0.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.0.ffn.0.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.0.ffn.2.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.0.ffn.2.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.0.modulation": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.0.norm3.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.0.norm3.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.0.self_attn.k.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.0.self_attn.k.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.0.self_attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.0.self_attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.0.self_attn.o.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.0.self_attn.o.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.0.self_attn.q.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.0.self_attn.q.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.0.self_attn.v.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.0.self_attn.v.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.1.cross_attn.k.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.1.cross_attn.k.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.1.cross_attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.1.cross_attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.1.cross_attn.o.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.1.cross_attn.o.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.1.cross_attn.q.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.1.cross_attn.q.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.1.cross_attn.v.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.1.cross_attn.v.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.1.ffn.0.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.1.ffn.0.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.1.ffn.2.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.1.ffn.2.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.1.modulation": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.1.norm3.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.1.norm3.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.1.self_attn.k.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.1.self_attn.k.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.1.self_attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.1.self_attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.1.self_attn.o.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.1.self_attn.o.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.1.self_attn.q.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.1.self_attn.q.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.1.self_attn.v.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.1.self_attn.v.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.10.cross_attn.k.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.10.cross_attn.k.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.10.cross_attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.10.cross_attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.10.cross_attn.o.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.10.cross_attn.o.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.10.cross_attn.q.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.10.cross_attn.q.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.10.cross_attn.v.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.10.cross_attn.v.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.10.ffn.0.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.10.ffn.0.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.10.ffn.2.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.10.ffn.2.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.10.modulation": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.10.norm3.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.10.norm3.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.10.self_attn.k.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.10.self_attn.k.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.10.self_attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.10.self_attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.10.self_attn.o.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.10.self_attn.o.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.10.self_attn.q.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.10.self_attn.q.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.10.self_attn.v.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.10.self_attn.v.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.11.cross_attn.k.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.11.cross_attn.k.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.11.cross_attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.11.cross_attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.11.cross_attn.o.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.11.cross_attn.o.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.11.cross_attn.q.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.11.cross_attn.q.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.11.cross_attn.v.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.11.cross_attn.v.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.11.ffn.0.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.11.ffn.0.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.11.ffn.2.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.11.ffn.2.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.11.modulation": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.11.norm3.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.11.norm3.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.11.self_attn.k.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.11.self_attn.k.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.11.self_attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.11.self_attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.11.self_attn.o.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.11.self_attn.o.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.11.self_attn.q.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.11.self_attn.q.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.11.self_attn.v.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.11.self_attn.v.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.12.cross_attn.k.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.12.cross_attn.k.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.12.cross_attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.12.cross_attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.12.cross_attn.o.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.12.cross_attn.o.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.12.cross_attn.q.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.12.cross_attn.q.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.12.cross_attn.v.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.12.cross_attn.v.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.12.ffn.0.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.12.ffn.0.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.12.ffn.2.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.12.ffn.2.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.12.modulation": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.12.norm3.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.12.norm3.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.12.self_attn.k.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.12.self_attn.k.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.12.self_attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.12.self_attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.12.self_attn.o.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.12.self_attn.o.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.12.self_attn.q.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.12.self_attn.q.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.12.self_attn.v.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.12.self_attn.v.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.13.cross_attn.k.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.13.cross_attn.k.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.13.cross_attn.norm_k.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.13.cross_attn.norm_q.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.13.cross_attn.o.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.13.cross_attn.o.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.13.cross_attn.q.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.13.cross_attn.q.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.13.cross_attn.v.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.13.cross_attn.v.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.13.ffn.0.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.13.ffn.0.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.13.ffn.2.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.13.ffn.2.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.13.modulation": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.13.norm3.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.13.norm3.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.13.self_attn.k.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.13.self_attn.k.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.13.self_attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.13.self_attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.13.self_attn.o.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.13.self_attn.o.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.13.self_attn.q.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.13.self_attn.q.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.13.self_attn.v.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.13.self_attn.v.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.14.cross_attn.k.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.14.cross_attn.k.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.14.cross_attn.norm_k.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.14.cross_attn.norm_q.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.14.cross_attn.o.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.14.cross_attn.o.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.14.cross_attn.q.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.14.cross_attn.q.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.14.cross_attn.v.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.14.cross_attn.v.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.14.ffn.0.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.14.ffn.0.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.14.ffn.2.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.14.ffn.2.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.14.modulation": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.14.norm3.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.14.norm3.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.14.self_attn.k.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.14.self_attn.k.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.14.self_attn.norm_k.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.14.self_attn.norm_q.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.14.self_attn.o.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.14.self_attn.o.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.14.self_attn.q.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.14.self_attn.q.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.14.self_attn.v.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.14.self_attn.v.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.15.cross_attn.k.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.15.cross_attn.k.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.15.cross_attn.norm_k.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.15.cross_attn.norm_q.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.15.cross_attn.o.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.15.cross_attn.o.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.15.cross_attn.q.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.15.cross_attn.q.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.15.cross_attn.v.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.15.cross_attn.v.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.15.ffn.0.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.15.ffn.0.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.15.ffn.2.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.15.ffn.2.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.15.modulation": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.15.norm3.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.15.norm3.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.15.self_attn.k.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.15.self_attn.k.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.15.self_attn.norm_k.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.15.self_attn.norm_q.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.15.self_attn.o.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.15.self_attn.o.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.15.self_attn.q.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.15.self_attn.q.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.15.self_attn.v.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.15.self_attn.v.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.16.cross_attn.k.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.16.cross_attn.k.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.16.cross_attn.norm_k.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.16.cross_attn.norm_q.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.16.cross_attn.o.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.16.cross_attn.o.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.16.cross_attn.q.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.16.cross_attn.q.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.16.cross_attn.v.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.16.cross_attn.v.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.16.ffn.0.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.16.ffn.0.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.16.ffn.2.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.16.ffn.2.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.16.modulation": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.16.norm3.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.16.norm3.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.16.self_attn.k.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.16.self_attn.k.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.16.self_attn.norm_k.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.16.self_attn.norm_q.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.16.self_attn.o.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.16.self_attn.o.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.16.self_attn.q.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.16.self_attn.q.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.16.self_attn.v.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.16.self_attn.v.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.17.cross_attn.k.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.17.cross_attn.k.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.17.cross_attn.norm_k.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.17.cross_attn.norm_q.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.17.cross_attn.o.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.17.cross_attn.o.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.17.cross_attn.q.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.17.cross_attn.q.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.17.cross_attn.v.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.17.cross_attn.v.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.17.ffn.0.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.17.ffn.0.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.17.ffn.2.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.17.ffn.2.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.17.modulation": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.17.norm3.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.17.norm3.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.17.self_attn.k.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.17.self_attn.k.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.17.self_attn.norm_k.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.17.self_attn.norm_q.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.17.self_attn.o.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.17.self_attn.o.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.17.self_attn.q.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.17.self_attn.q.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.17.self_attn.v.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.17.self_attn.v.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.18.cross_attn.k.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.18.cross_attn.k.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.18.cross_attn.norm_k.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.18.cross_attn.norm_q.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.18.cross_attn.o.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.18.cross_attn.o.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.18.cross_attn.q.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.18.cross_attn.q.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.18.cross_attn.v.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.18.cross_attn.v.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.18.ffn.0.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.18.ffn.0.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.18.ffn.2.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.18.ffn.2.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.18.modulation": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.18.norm3.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.18.norm3.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.18.self_attn.k.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.18.self_attn.k.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.18.self_attn.norm_k.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.18.self_attn.norm_q.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.18.self_attn.o.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.18.self_attn.o.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.18.self_attn.q.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.18.self_attn.q.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.18.self_attn.v.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.18.self_attn.v.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.19.cross_attn.k.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.19.cross_attn.k.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.19.cross_attn.norm_k.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.19.cross_attn.norm_q.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.19.cross_attn.o.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.19.cross_attn.o.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.19.cross_attn.q.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.19.cross_attn.q.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.19.cross_attn.v.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.19.cross_attn.v.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.19.ffn.0.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.19.ffn.0.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.19.ffn.2.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.19.ffn.2.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.19.modulation": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.19.norm3.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.19.norm3.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.19.self_attn.k.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.19.self_attn.k.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.19.self_attn.norm_k.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.19.self_attn.norm_q.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.19.self_attn.o.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.19.self_attn.o.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.19.self_attn.q.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.19.self_attn.q.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.19.self_attn.v.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.19.self_attn.v.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.2.cross_attn.k.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.2.cross_attn.k.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.2.cross_attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.2.cross_attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.2.cross_attn.o.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.2.cross_attn.o.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.2.cross_attn.q.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.2.cross_attn.q.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.2.cross_attn.v.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.2.cross_attn.v.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.2.ffn.0.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.2.ffn.0.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.2.ffn.2.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.2.ffn.2.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.2.modulation": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.2.norm3.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.2.norm3.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.2.self_attn.k.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.2.self_attn.k.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.2.self_attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.2.self_attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.2.self_attn.o.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.2.self_attn.o.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.2.self_attn.q.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.2.self_attn.q.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.2.self_attn.v.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.2.self_attn.v.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.20.cross_attn.k.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.20.cross_attn.k.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.20.cross_attn.norm_k.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.20.cross_attn.norm_q.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.20.cross_attn.o.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.20.cross_attn.o.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.20.cross_attn.q.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.20.cross_attn.q.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.20.cross_attn.v.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.20.cross_attn.v.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.20.ffn.0.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.20.ffn.0.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.20.ffn.2.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.20.ffn.2.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.20.modulation": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.20.norm3.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.20.norm3.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.20.self_attn.k.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.20.self_attn.k.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.20.self_attn.norm_k.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.20.self_attn.norm_q.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.20.self_attn.o.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.20.self_attn.o.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.20.self_attn.q.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.20.self_attn.q.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.20.self_attn.v.bias": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.20.self_attn.v.weight": "diffusion_pytorch_model-00003-of-00006.safetensors",
+    "blocks.21.cross_attn.k.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.21.cross_attn.k.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.21.cross_attn.norm_k.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.21.cross_attn.norm_q.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.21.cross_attn.o.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.21.cross_attn.o.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.21.cross_attn.q.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.21.cross_attn.q.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.21.cross_attn.v.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.21.cross_attn.v.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.21.ffn.0.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.21.ffn.0.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.21.ffn.2.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.21.ffn.2.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.21.modulation": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.21.norm3.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.21.norm3.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.21.self_attn.k.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.21.self_attn.k.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.21.self_attn.norm_k.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.21.self_attn.norm_q.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.21.self_attn.o.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.21.self_attn.o.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.21.self_attn.q.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.21.self_attn.q.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.21.self_attn.v.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.21.self_attn.v.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.22.cross_attn.k.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.22.cross_attn.k.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.22.cross_attn.norm_k.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.22.cross_attn.norm_q.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.22.cross_attn.o.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.22.cross_attn.o.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.22.cross_attn.q.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.22.cross_attn.q.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.22.cross_attn.v.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.22.cross_attn.v.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.22.ffn.0.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.22.ffn.0.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.22.ffn.2.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.22.ffn.2.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.22.modulation": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.22.norm3.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.22.norm3.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.22.self_attn.k.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.22.self_attn.k.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.22.self_attn.norm_k.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.22.self_attn.norm_q.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.22.self_attn.o.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.22.self_attn.o.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.22.self_attn.q.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.22.self_attn.q.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.22.self_attn.v.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.22.self_attn.v.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.23.cross_attn.k.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.23.cross_attn.k.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.23.cross_attn.norm_k.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.23.cross_attn.norm_q.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.23.cross_attn.o.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.23.cross_attn.o.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.23.cross_attn.q.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.23.cross_attn.q.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.23.cross_attn.v.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.23.cross_attn.v.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.23.ffn.0.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.23.ffn.0.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.23.ffn.2.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.23.ffn.2.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.23.modulation": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.23.norm3.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.23.norm3.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.23.self_attn.k.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.23.self_attn.k.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.23.self_attn.norm_k.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.23.self_attn.norm_q.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.23.self_attn.o.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.23.self_attn.o.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.23.self_attn.q.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.23.self_attn.q.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.23.self_attn.v.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.23.self_attn.v.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.24.cross_attn.k.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.24.cross_attn.k.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.24.cross_attn.norm_k.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.24.cross_attn.norm_q.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.24.cross_attn.o.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.24.cross_attn.o.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.24.cross_attn.q.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.24.cross_attn.q.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.24.cross_attn.v.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.24.cross_attn.v.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.24.ffn.0.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.24.ffn.0.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.24.ffn.2.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.24.ffn.2.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.24.modulation": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.24.norm3.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.24.norm3.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.24.self_attn.k.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.24.self_attn.k.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.24.self_attn.norm_k.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.24.self_attn.norm_q.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.24.self_attn.o.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.24.self_attn.o.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.24.self_attn.q.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.24.self_attn.q.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.24.self_attn.v.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.24.self_attn.v.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.25.cross_attn.k.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.25.cross_attn.k.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.25.cross_attn.norm_k.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.25.cross_attn.norm_q.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.25.cross_attn.o.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.25.cross_attn.o.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.25.cross_attn.q.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.25.cross_attn.q.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.25.cross_attn.v.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.25.cross_attn.v.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.25.ffn.0.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.25.ffn.0.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.25.ffn.2.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.25.ffn.2.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.25.modulation": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.25.norm3.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.25.norm3.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.25.self_attn.k.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.25.self_attn.k.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.25.self_attn.norm_k.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.25.self_attn.norm_q.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.25.self_attn.o.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.25.self_attn.o.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.25.self_attn.q.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.25.self_attn.q.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.25.self_attn.v.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.25.self_attn.v.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.26.cross_attn.k.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.26.cross_attn.k.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.26.cross_attn.norm_k.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.26.cross_attn.norm_q.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.26.cross_attn.o.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.26.cross_attn.o.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.26.cross_attn.q.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.26.cross_attn.q.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.26.cross_attn.v.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.26.cross_attn.v.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.26.ffn.0.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.26.ffn.0.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.26.ffn.2.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.26.ffn.2.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.26.modulation": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.26.norm3.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.26.norm3.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.26.self_attn.k.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.26.self_attn.k.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.26.self_attn.norm_k.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.26.self_attn.norm_q.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.26.self_attn.o.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.26.self_attn.o.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.26.self_attn.q.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.26.self_attn.q.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.26.self_attn.v.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.26.self_attn.v.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.27.cross_attn.k.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.27.cross_attn.k.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.27.cross_attn.norm_k.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.27.cross_attn.norm_q.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.27.cross_attn.o.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.27.cross_attn.o.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.27.cross_attn.q.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.27.cross_attn.q.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.27.cross_attn.v.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.27.cross_attn.v.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.27.ffn.0.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.27.ffn.0.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.27.ffn.2.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.27.ffn.2.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.27.modulation": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.27.norm3.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.27.norm3.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.27.self_attn.k.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.27.self_attn.k.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.27.self_attn.norm_k.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.27.self_attn.norm_q.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.27.self_attn.o.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.27.self_attn.o.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.27.self_attn.q.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.27.self_attn.q.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.27.self_attn.v.bias": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.27.self_attn.v.weight": "diffusion_pytorch_model-00004-of-00006.safetensors",
+    "blocks.28.cross_attn.k.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.28.cross_attn.k.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.28.cross_attn.norm_k.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.28.cross_attn.norm_q.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.28.cross_attn.o.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.28.cross_attn.o.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.28.cross_attn.q.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.28.cross_attn.q.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.28.cross_attn.v.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.28.cross_attn.v.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.28.ffn.0.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.28.ffn.0.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.28.ffn.2.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.28.ffn.2.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.28.modulation": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.28.norm3.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.28.norm3.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.28.self_attn.k.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.28.self_attn.k.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.28.self_attn.norm_k.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.28.self_attn.norm_q.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.28.self_attn.o.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.28.self_attn.o.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.28.self_attn.q.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.28.self_attn.q.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.28.self_attn.v.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.28.self_attn.v.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.29.cross_attn.k.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.29.cross_attn.k.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.29.cross_attn.norm_k.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.29.cross_attn.norm_q.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.29.cross_attn.o.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.29.cross_attn.o.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.29.cross_attn.q.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.29.cross_attn.q.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.29.cross_attn.v.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.29.cross_attn.v.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.29.ffn.0.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.29.ffn.0.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.29.ffn.2.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.29.ffn.2.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.29.modulation": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.29.norm3.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.29.norm3.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.29.self_attn.k.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.29.self_attn.k.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.29.self_attn.norm_k.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.29.self_attn.norm_q.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.29.self_attn.o.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.29.self_attn.o.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.29.self_attn.q.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.29.self_attn.q.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.29.self_attn.v.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.29.self_attn.v.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.3.cross_attn.k.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.3.cross_attn.k.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.3.cross_attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.3.cross_attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.3.cross_attn.o.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.3.cross_attn.o.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.3.cross_attn.q.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.3.cross_attn.q.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.3.cross_attn.v.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.3.cross_attn.v.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.3.ffn.0.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.3.ffn.0.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.3.ffn.2.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.3.ffn.2.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.3.modulation": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.3.norm3.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.3.norm3.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.3.self_attn.k.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.3.self_attn.k.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.3.self_attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.3.self_attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.3.self_attn.o.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.3.self_attn.o.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.3.self_attn.q.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.3.self_attn.q.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.3.self_attn.v.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.3.self_attn.v.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.30.cross_attn.k.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.30.cross_attn.k.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.30.cross_attn.norm_k.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.30.cross_attn.norm_q.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.30.cross_attn.o.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.30.cross_attn.o.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.30.cross_attn.q.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.30.cross_attn.q.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.30.cross_attn.v.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.30.cross_attn.v.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.30.ffn.0.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.30.ffn.0.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.30.ffn.2.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.30.ffn.2.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.30.modulation": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.30.norm3.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.30.norm3.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.30.self_attn.k.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.30.self_attn.k.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.30.self_attn.norm_k.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.30.self_attn.norm_q.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.30.self_attn.o.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.30.self_attn.o.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.30.self_attn.q.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.30.self_attn.q.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.30.self_attn.v.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.30.self_attn.v.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.31.cross_attn.k.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.31.cross_attn.k.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.31.cross_attn.norm_k.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.31.cross_attn.norm_q.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.31.cross_attn.o.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.31.cross_attn.o.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.31.cross_attn.q.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.31.cross_attn.q.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.31.cross_attn.v.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.31.cross_attn.v.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.31.ffn.0.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.31.ffn.0.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.31.ffn.2.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.31.ffn.2.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.31.modulation": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.31.norm3.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.31.norm3.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.31.self_attn.k.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.31.self_attn.k.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.31.self_attn.norm_k.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.31.self_attn.norm_q.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.31.self_attn.o.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.31.self_attn.o.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.31.self_attn.q.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.31.self_attn.q.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.31.self_attn.v.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.31.self_attn.v.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.32.cross_attn.k.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.32.cross_attn.k.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.32.cross_attn.norm_k.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.32.cross_attn.norm_q.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.32.cross_attn.o.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.32.cross_attn.o.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.32.cross_attn.q.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.32.cross_attn.q.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.32.cross_attn.v.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.32.cross_attn.v.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.32.ffn.0.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.32.ffn.0.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.32.ffn.2.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.32.ffn.2.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.32.modulation": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.32.norm3.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.32.norm3.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.32.self_attn.k.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.32.self_attn.k.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.32.self_attn.norm_k.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.32.self_attn.norm_q.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.32.self_attn.o.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.32.self_attn.o.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.32.self_attn.q.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.32.self_attn.q.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.32.self_attn.v.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.32.self_attn.v.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.33.cross_attn.k.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.33.cross_attn.k.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.33.cross_attn.norm_k.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.33.cross_attn.norm_q.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.33.cross_attn.o.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.33.cross_attn.o.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.33.cross_attn.q.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.33.cross_attn.q.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.33.cross_attn.v.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.33.cross_attn.v.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.33.ffn.0.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.33.ffn.0.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.33.ffn.2.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.33.ffn.2.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.33.modulation": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.33.norm3.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.33.norm3.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.33.self_attn.k.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.33.self_attn.k.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.33.self_attn.norm_k.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.33.self_attn.norm_q.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.33.self_attn.o.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.33.self_attn.o.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.33.self_attn.q.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.33.self_attn.q.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.33.self_attn.v.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.33.self_attn.v.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.34.cross_attn.k.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.34.cross_attn.k.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.34.cross_attn.norm_k.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.34.cross_attn.norm_q.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.34.cross_attn.o.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.34.cross_attn.o.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.34.cross_attn.q.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.34.cross_attn.q.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.34.cross_attn.v.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.34.cross_attn.v.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.34.ffn.0.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.34.ffn.0.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.34.ffn.2.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.34.ffn.2.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.34.modulation": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.34.norm3.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.34.norm3.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.34.self_attn.k.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.34.self_attn.k.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.34.self_attn.norm_k.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.34.self_attn.norm_q.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.34.self_attn.o.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.34.self_attn.o.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.34.self_attn.q.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.34.self_attn.q.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.34.self_attn.v.bias": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.34.self_attn.v.weight": "diffusion_pytorch_model-00005-of-00006.safetensors",
+    "blocks.35.cross_attn.k.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.35.cross_attn.k.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.35.cross_attn.norm_k.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.35.cross_attn.norm_q.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.35.cross_attn.o.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.35.cross_attn.o.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.35.cross_attn.q.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.35.cross_attn.q.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.35.cross_attn.v.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.35.cross_attn.v.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.35.ffn.0.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.35.ffn.0.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.35.ffn.2.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.35.ffn.2.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.35.modulation": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.35.norm3.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.35.norm3.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.35.self_attn.k.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.35.self_attn.k.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.35.self_attn.norm_k.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.35.self_attn.norm_q.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.35.self_attn.o.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.35.self_attn.o.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.35.self_attn.q.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.35.self_attn.q.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.35.self_attn.v.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.35.self_attn.v.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.36.cross_attn.k.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.36.cross_attn.k.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.36.cross_attn.norm_k.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.36.cross_attn.norm_q.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.36.cross_attn.o.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.36.cross_attn.o.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.36.cross_attn.q.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.36.cross_attn.q.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.36.cross_attn.v.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.36.cross_attn.v.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.36.ffn.0.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.36.ffn.0.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.36.ffn.2.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.36.ffn.2.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.36.modulation": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.36.norm3.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.36.norm3.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.36.self_attn.k.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.36.self_attn.k.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.36.self_attn.norm_k.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.36.self_attn.norm_q.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.36.self_attn.o.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.36.self_attn.o.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.36.self_attn.q.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.36.self_attn.q.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.36.self_attn.v.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.36.self_attn.v.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.37.cross_attn.k.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.37.cross_attn.k.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.37.cross_attn.norm_k.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.37.cross_attn.norm_q.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.37.cross_attn.o.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.37.cross_attn.o.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.37.cross_attn.q.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.37.cross_attn.q.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.37.cross_attn.v.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.37.cross_attn.v.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.37.ffn.0.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.37.ffn.0.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.37.ffn.2.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.37.ffn.2.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.37.modulation": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.37.norm3.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.37.norm3.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.37.self_attn.k.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.37.self_attn.k.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.37.self_attn.norm_k.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.37.self_attn.norm_q.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.37.self_attn.o.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.37.self_attn.o.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.37.self_attn.q.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.37.self_attn.q.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.37.self_attn.v.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.37.self_attn.v.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.38.cross_attn.k.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.38.cross_attn.k.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.38.cross_attn.norm_k.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.38.cross_attn.norm_q.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.38.cross_attn.o.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.38.cross_attn.o.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.38.cross_attn.q.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.38.cross_attn.q.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.38.cross_attn.v.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.38.cross_attn.v.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.38.ffn.0.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.38.ffn.0.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.38.ffn.2.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.38.ffn.2.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.38.modulation": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.38.norm3.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.38.norm3.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.38.self_attn.k.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.38.self_attn.k.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.38.self_attn.norm_k.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.38.self_attn.norm_q.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.38.self_attn.o.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.38.self_attn.o.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.38.self_attn.q.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.38.self_attn.q.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.38.self_attn.v.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.38.self_attn.v.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.39.cross_attn.k.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.39.cross_attn.k.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.39.cross_attn.norm_k.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.39.cross_attn.norm_q.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.39.cross_attn.o.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.39.cross_attn.o.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.39.cross_attn.q.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.39.cross_attn.q.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.39.cross_attn.v.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.39.cross_attn.v.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.39.ffn.0.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.39.ffn.0.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.39.ffn.2.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.39.ffn.2.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.39.modulation": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.39.norm3.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.39.norm3.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.39.self_attn.k.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.39.self_attn.k.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.39.self_attn.norm_k.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.39.self_attn.norm_q.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.39.self_attn.o.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.39.self_attn.o.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.39.self_attn.q.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.39.self_attn.q.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.39.self_attn.v.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.39.self_attn.v.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "blocks.4.cross_attn.k.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.4.cross_attn.k.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.4.cross_attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.4.cross_attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.4.cross_attn.o.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.4.cross_attn.o.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.4.cross_attn.q.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.4.cross_attn.q.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.4.cross_attn.v.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.4.cross_attn.v.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.4.ffn.0.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.4.ffn.0.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.4.ffn.2.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.4.ffn.2.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.4.modulation": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.4.norm3.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.4.norm3.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.4.self_attn.k.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.4.self_attn.k.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.4.self_attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.4.self_attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.4.self_attn.o.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.4.self_attn.o.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.4.self_attn.q.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.4.self_attn.q.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.4.self_attn.v.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.4.self_attn.v.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.5.cross_attn.k.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.5.cross_attn.k.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.5.cross_attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.5.cross_attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.5.cross_attn.o.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.5.cross_attn.o.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.5.cross_attn.q.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.5.cross_attn.q.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.5.cross_attn.v.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.5.cross_attn.v.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.5.ffn.0.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.5.ffn.0.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.5.ffn.2.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.5.ffn.2.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.5.modulation": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.5.norm3.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.5.norm3.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.5.self_attn.k.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.5.self_attn.k.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.5.self_attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.5.self_attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.5.self_attn.o.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.5.self_attn.o.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.5.self_attn.q.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.5.self_attn.q.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.5.self_attn.v.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.5.self_attn.v.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.6.cross_attn.k.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.6.cross_attn.k.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.6.cross_attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.6.cross_attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.6.cross_attn.o.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.6.cross_attn.o.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.6.cross_attn.q.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.6.cross_attn.q.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.6.cross_attn.v.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.6.cross_attn.v.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.6.ffn.0.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.6.ffn.0.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.6.ffn.2.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.6.ffn.2.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.6.modulation": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.6.norm3.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.6.norm3.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.6.self_attn.k.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.6.self_attn.k.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.6.self_attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.6.self_attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.6.self_attn.o.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.6.self_attn.o.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.6.self_attn.q.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.6.self_attn.q.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.6.self_attn.v.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.6.self_attn.v.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "blocks.7.cross_attn.k.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.7.cross_attn.k.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.7.cross_attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.7.cross_attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.7.cross_attn.o.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.7.cross_attn.o.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.7.cross_attn.q.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.7.cross_attn.q.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.7.cross_attn.v.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.7.cross_attn.v.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.7.ffn.0.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.7.ffn.0.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.7.ffn.2.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.7.ffn.2.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.7.modulation": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.7.norm3.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.7.norm3.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.7.self_attn.k.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.7.self_attn.k.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.7.self_attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.7.self_attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.7.self_attn.o.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.7.self_attn.o.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.7.self_attn.q.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.7.self_attn.q.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.7.self_attn.v.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.7.self_attn.v.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.8.cross_attn.k.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.8.cross_attn.k.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.8.cross_attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.8.cross_attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.8.cross_attn.o.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.8.cross_attn.o.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.8.cross_attn.q.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.8.cross_attn.q.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.8.cross_attn.v.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.8.cross_attn.v.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.8.ffn.0.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.8.ffn.0.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.8.ffn.2.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.8.ffn.2.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.8.modulation": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.8.norm3.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.8.norm3.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.8.self_attn.k.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.8.self_attn.k.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.8.self_attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.8.self_attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.8.self_attn.o.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.8.self_attn.o.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.8.self_attn.q.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.8.self_attn.q.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.8.self_attn.v.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.8.self_attn.v.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.9.cross_attn.k.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.9.cross_attn.k.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.9.cross_attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.9.cross_attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.9.cross_attn.o.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.9.cross_attn.o.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.9.cross_attn.q.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.9.cross_attn.q.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.9.cross_attn.v.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.9.cross_attn.v.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.9.ffn.0.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.9.ffn.0.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.9.ffn.2.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.9.ffn.2.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.9.modulation": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.9.norm3.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.9.norm3.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.9.self_attn.k.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.9.self_attn.k.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.9.self_attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.9.self_attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.9.self_attn.o.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.9.self_attn.o.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.9.self_attn.q.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.9.self_attn.q.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.9.self_attn.v.bias": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "blocks.9.self_attn.v.weight": "diffusion_pytorch_model-00002-of-00006.safetensors",
+    "head.head.bias": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "head.head.weight": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "head.modulation": "diffusion_pytorch_model-00006-of-00006.safetensors",
+    "patch_embedding.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "patch_embedding.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "text_embedding.0.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "text_embedding.0.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "text_embedding.2.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "text_embedding.2.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "time_embedding.0.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "time_embedding.0.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "time_embedding.2.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "time_embedding.2.weight": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "time_projection.1.bias": "diffusion_pytorch_model-00001-of-00006.safetensors",
+    "time_projection.1.weight": "diffusion_pytorch_model-00001-of-00006.safetensors"
+  }
+}

wan_models/Wan2.1-T2V-14B/google/umt5-xxl/special_tokens_map.json

@@ -0,0 +1,308 @@
+{
+  "additional_special_tokens": [
+    "<extra_id_0>",
+    "<extra_id_1>",
+    "<extra_id_2>",
+    "<extra_id_3>",
+    "<extra_id_4>",
+    "<extra_id_5>",
+    "<extra_id_6>",
+    "<extra_id_7>",
+    "<extra_id_8>",
+    "<extra_id_9>",
+    "<extra_id_10>",
+    "<extra_id_11>",
+    "<extra_id_12>",
+    "<extra_id_13>",
+    "<extra_id_14>",
+    "<extra_id_15>",
+    "<extra_id_16>",
+    "<extra_id_17>",
+    "<extra_id_18>",
+    "<extra_id_19>",
+    "<extra_id_20>",
+    "<extra_id_21>",
+    "<extra_id_22>",
+    "<extra_id_23>",
+    "<extra_id_24>",
+    "<extra_id_25>",
+    "<extra_id_26>",
+    "<extra_id_27>",
+    "<extra_id_28>",
+    "<extra_id_29>",
+    "<extra_id_30>",
+    "<extra_id_31>",
+    "<extra_id_32>",
+    "<extra_id_33>",
+    "<extra_id_34>",
+    "<extra_id_35>",
+    "<extra_id_36>",
+    "<extra_id_37>",
+    "<extra_id_38>",
+    "<extra_id_39>",
+    "<extra_id_40>",
+    "<extra_id_41>",
+    "<extra_id_42>",
+    "<extra_id_43>",
+    "<extra_id_44>",
+    "<extra_id_45>",
+    "<extra_id_46>",
+    "<extra_id_47>",
+    "<extra_id_48>",
+    "<extra_id_49>",
+    "<extra_id_50>",
+    "<extra_id_51>",
+    "<extra_id_52>",
+    "<extra_id_53>",
+    "<extra_id_54>",
+    "<extra_id_55>",
+    "<extra_id_56>",
+    "<extra_id_57>",
+    "<extra_id_58>",
+    "<extra_id_59>",
+    "<extra_id_60>",
+    "<extra_id_61>",
+    "<extra_id_62>",
+    "<extra_id_63>",
+    "<extra_id_64>",
+    "<extra_id_65>",
+    "<extra_id_66>",
+    "<extra_id_67>",
+    "<extra_id_68>",
+    "<extra_id_69>",
+    "<extra_id_70>",
+    "<extra_id_71>",
+    "<extra_id_72>",
+    "<extra_id_73>",
+    "<extra_id_74>",
+    "<extra_id_75>",
+    "<extra_id_76>",
+    "<extra_id_77>",
+    "<extra_id_78>",
+    "<extra_id_79>",
+    "<extra_id_80>",
+    "<extra_id_81>",
+    "<extra_id_82>",
+    "<extra_id_83>",
+    "<extra_id_84>",
+    "<extra_id_85>",
+    "<extra_id_86>",
+    "<extra_id_87>",
+    "<extra_id_88>",
+    "<extra_id_89>",
+    "<extra_id_90>",
+    "<extra_id_91>",
+    "<extra_id_92>",
+    "<extra_id_93>",
+    "<extra_id_94>",
+    "<extra_id_95>",
+    "<extra_id_96>",
+    "<extra_id_97>",
+    "<extra_id_98>",
+    "<extra_id_99>",
+    "<extra_id_100>",
+    "<extra_id_101>",
+    "<extra_id_102>",
+    "<extra_id_103>",
+    "<extra_id_104>",
+    "<extra_id_105>",
+    "<extra_id_106>",
+    "<extra_id_107>",
+    "<extra_id_108>",
+    "<extra_id_109>",
+    "<extra_id_110>",
+    "<extra_id_111>",
+    "<extra_id_112>",
+    "<extra_id_113>",
+    "<extra_id_114>",
+    "<extra_id_115>",
+    "<extra_id_116>",
+    "<extra_id_117>",
+    "<extra_id_118>",
+    "<extra_id_119>",
+    "<extra_id_120>",
+    "<extra_id_121>",
+    "<extra_id_122>",
+    "<extra_id_123>",
+    "<extra_id_124>",
+    "<extra_id_125>",
+    "<extra_id_126>",
+    "<extra_id_127>",
+    "<extra_id_128>",
+    "<extra_id_129>",
+    "<extra_id_130>",
+    "<extra_id_131>",
+    "<extra_id_132>",
+    "<extra_id_133>",
+    "<extra_id_134>",
+    "<extra_id_135>",
+    "<extra_id_136>",
+    "<extra_id_137>",
+    "<extra_id_138>",
+    "<extra_id_139>",
+    "<extra_id_140>",
+    "<extra_id_141>",
+    "<extra_id_142>",
+    "<extra_id_143>",
+    "<extra_id_144>",
+    "<extra_id_145>",
+    "<extra_id_146>",
+    "<extra_id_147>",
+    "<extra_id_148>",
+    "<extra_id_149>",
+    "<extra_id_150>",
+    "<extra_id_151>",
+    "<extra_id_152>",
+    "<extra_id_153>",
+    "<extra_id_154>",
+    "<extra_id_155>",
+    "<extra_id_156>",
+    "<extra_id_157>",
+    "<extra_id_158>",
+    "<extra_id_159>",
+    "<extra_id_160>",
+    "<extra_id_161>",
+    "<extra_id_162>",
+    "<extra_id_163>",
+    "<extra_id_164>",
+    "<extra_id_165>",
+    "<extra_id_166>",
+    "<extra_id_167>",
+    "<extra_id_168>",
+    "<extra_id_169>",
+    "<extra_id_170>",
+    "<extra_id_171>",
+    "<extra_id_172>",
+    "<extra_id_173>",
+    "<extra_id_174>",
+    "<extra_id_175>",
+    "<extra_id_176>",
+    "<extra_id_177>",
+    "<extra_id_178>",
+    "<extra_id_179>",
+    "<extra_id_180>",
+    "<extra_id_181>",
+    "<extra_id_182>",
+    "<extra_id_183>",
+    "<extra_id_184>",
+    "<extra_id_185>",
+    "<extra_id_186>",
+    "<extra_id_187>",
+    "<extra_id_188>",
+    "<extra_id_189>",
+    "<extra_id_190>",
+    "<extra_id_191>",
+    "<extra_id_192>",
+    "<extra_id_193>",
+    "<extra_id_194>",
+    "<extra_id_195>",
+    "<extra_id_196>",
+    "<extra_id_197>",
+    "<extra_id_198>",
+    "<extra_id_199>",
+    "<extra_id_200>",
+    "<extra_id_201>",
+    "<extra_id_202>",
+    "<extra_id_203>",
+    "<extra_id_204>",
+    "<extra_id_205>",
+    "<extra_id_206>",
+    "<extra_id_207>",
+    "<extra_id_208>",
+    "<extra_id_209>",
+    "<extra_id_210>",
+    "<extra_id_211>",
+    "<extra_id_212>",
+    "<extra_id_213>",
+    "<extra_id_214>",
+    "<extra_id_215>",
+    "<extra_id_216>",
+    "<extra_id_217>",
+    "<extra_id_218>",
+    "<extra_id_219>",
+    "<extra_id_220>",
+    "<extra_id_221>",
+    "<extra_id_222>",
+    "<extra_id_223>",
+    "<extra_id_224>",
+    "<extra_id_225>",
+    "<extra_id_226>",
+    "<extra_id_227>",
+    "<extra_id_228>",
+    "<extra_id_229>",
+    "<extra_id_230>",
+    "<extra_id_231>",
+    "<extra_id_232>",
+    "<extra_id_233>",
+    "<extra_id_234>",
+    "<extra_id_235>",
+    "<extra_id_236>",
+    "<extra_id_237>",
+    "<extra_id_238>",
+    "<extra_id_239>",
+    "<extra_id_240>",
+    "<extra_id_241>",
+    "<extra_id_242>",
+    "<extra_id_243>",
+    "<extra_id_244>",
+    "<extra_id_245>",
+    "<extra_id_246>",
+    "<extra_id_247>",
+    "<extra_id_248>",
+    "<extra_id_249>",
+    "<extra_id_250>",
+    "<extra_id_251>",
+    "<extra_id_252>",
+    "<extra_id_253>",
+    "<extra_id_254>",
+    "<extra_id_255>",
+    "<extra_id_256>",
+    "<extra_id_257>",
+    "<extra_id_258>",
+    "<extra_id_259>",
+    "<extra_id_260>",
+    "<extra_id_261>",
+    "<extra_id_262>",
+    "<extra_id_263>",
+    "<extra_id_264>",
+    "<extra_id_265>",
+    "<extra_id_266>",
+    "<extra_id_267>",
+    "<extra_id_268>",
+    "<extra_id_269>",
+    "<extra_id_270>",
+    "<extra_id_271>",
+    "<extra_id_272>",
+    "<extra_id_273>",
+    "<extra_id_274>",
+    "<extra_id_275>",
+    "<extra_id_276>",
+    "<extra_id_277>",
+    "<extra_id_278>",
+    "<extra_id_279>",
+    "<extra_id_280>",
+    "<extra_id_281>",
+    "<extra_id_282>",
+    "<extra_id_283>",
+    "<extra_id_284>",
+    "<extra_id_285>",
+    "<extra_id_286>",
+    "<extra_id_287>",
+    "<extra_id_288>",
+    "<extra_id_289>",
+    "<extra_id_290>",
+    "<extra_id_291>",
+    "<extra_id_292>",
+    "<extra_id_293>",
+    "<extra_id_294>",
+    "<extra_id_295>",
+    "<extra_id_296>",
+    "<extra_id_297>",
+    "<extra_id_298>",
+    "<extra_id_299>"
+  ],
+  "bos_token": "<s>",
+  "eos_token": "</s>",
+  "pad_token": "<pad>",
+  "unk_token": "<unk>"
+}

wan_models/Wan2.1-T2V-14B/google/umt5-xxl/tokenizer_config.json

@@ -0,0 +1,2748 @@
+{
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<pad>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "</s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "<s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "3": {
+      "content": "<unk>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256000": {
+      "content": "<extra_id_299>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256001": {
+      "content": "<extra_id_298>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256002": {
+      "content": "<extra_id_297>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256003": {
+      "content": "<extra_id_296>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256004": {
+      "content": "<extra_id_295>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256005": {
+      "content": "<extra_id_294>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256006": {
+      "content": "<extra_id_293>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256007": {
+      "content": "<extra_id_292>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256008": {
+      "content": "<extra_id_291>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256009": {
+      "content": "<extra_id_290>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256010": {
+      "content": "<extra_id_289>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256011": {
+      "content": "<extra_id_288>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256012": {
+      "content": "<extra_id_287>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256013": {
+      "content": "<extra_id_286>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256014": {
+      "content": "<extra_id_285>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256015": {
+      "content": "<extra_id_284>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256016": {
+      "content": "<extra_id_283>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256017": {
+      "content": "<extra_id_282>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256018": {
+      "content": "<extra_id_281>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256019": {
+      "content": "<extra_id_280>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256020": {
+      "content": "<extra_id_279>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256021": {
+      "content": "<extra_id_278>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256022": {
+      "content": "<extra_id_277>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256023": {
+      "content": "<extra_id_276>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256024": {
+      "content": "<extra_id_275>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256025": {
+      "content": "<extra_id_274>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256026": {
+      "content": "<extra_id_273>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256027": {
+      "content": "<extra_id_272>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256028": {
+      "content": "<extra_id_271>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256029": {
+      "content": "<extra_id_270>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256030": {
+      "content": "<extra_id_269>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256031": {
+      "content": "<extra_id_268>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256032": {
+      "content": "<extra_id_267>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256033": {
+      "content": "<extra_id_266>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256034": {
+      "content": "<extra_id_265>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256035": {
+      "content": "<extra_id_264>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256036": {
+      "content": "<extra_id_263>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256037": {
+      "content": "<extra_id_262>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256038": {
+      "content": "<extra_id_261>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256039": {
+      "content": "<extra_id_260>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256040": {
+      "content": "<extra_id_259>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256041": {
+      "content": "<extra_id_258>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256042": {
+      "content": "<extra_id_257>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256043": {
+      "content": "<extra_id_256>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256044": {
+      "content": "<extra_id_255>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256045": {
+      "content": "<extra_id_254>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256046": {
+      "content": "<extra_id_253>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256047": {
+      "content": "<extra_id_252>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256048": {
+      "content": "<extra_id_251>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256049": {
+      "content": "<extra_id_250>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256050": {
+      "content": "<extra_id_249>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256051": {
+      "content": "<extra_id_248>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256052": {
+      "content": "<extra_id_247>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256053": {
+      "content": "<extra_id_246>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256054": {
+      "content": "<extra_id_245>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256055": {
+      "content": "<extra_id_244>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256056": {
+      "content": "<extra_id_243>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256057": {
+      "content": "<extra_id_242>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256058": {
+      "content": "<extra_id_241>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256059": {
+      "content": "<extra_id_240>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256060": {
+      "content": "<extra_id_239>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256061": {
+      "content": "<extra_id_238>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256062": {
+      "content": "<extra_id_237>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256063": {
+      "content": "<extra_id_236>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256064": {
+      "content": "<extra_id_235>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256065": {
+      "content": "<extra_id_234>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256066": {
+      "content": "<extra_id_233>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256067": {
+      "content": "<extra_id_232>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256068": {
+      "content": "<extra_id_231>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256069": {
+      "content": "<extra_id_230>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256070": {
+      "content": "<extra_id_229>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256071": {
+      "content": "<extra_id_228>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256072": {
+      "content": "<extra_id_227>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256073": {
+      "content": "<extra_id_226>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256074": {
+      "content": "<extra_id_225>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256075": {
+      "content": "<extra_id_224>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256076": {
+      "content": "<extra_id_223>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256077": {
+      "content": "<extra_id_222>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256078": {
+      "content": "<extra_id_221>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256079": {
+      "content": "<extra_id_220>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256080": {
+      "content": "<extra_id_219>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256081": {
+      "content": "<extra_id_218>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256082": {
+      "content": "<extra_id_217>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256083": {
+      "content": "<extra_id_216>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256084": {
+      "content": "<extra_id_215>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256085": {
+      "content": "<extra_id_214>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256086": {
+      "content": "<extra_id_213>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256087": {
+      "content": "<extra_id_212>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256088": {
+      "content": "<extra_id_211>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256089": {
+      "content": "<extra_id_210>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256090": {
+      "content": "<extra_id_209>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256091": {
+      "content": "<extra_id_208>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256092": {
+      "content": "<extra_id_207>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256093": {
+      "content": "<extra_id_206>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256094": {
+      "content": "<extra_id_205>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256095": {
+      "content": "<extra_id_204>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256096": {
+      "content": "<extra_id_203>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256097": {
+      "content": "<extra_id_202>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256098": {
+      "content": "<extra_id_201>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256099": {
+      "content": "<extra_id_200>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256100": {
+      "content": "<extra_id_199>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256101": {
+      "content": "<extra_id_198>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256102": {
+      "content": "<extra_id_197>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256103": {
+      "content": "<extra_id_196>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256104": {
+      "content": "<extra_id_195>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256105": {
+      "content": "<extra_id_194>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256106": {
+      "content": "<extra_id_193>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256107": {
+      "content": "<extra_id_192>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256108": {
+      "content": "<extra_id_191>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256109": {
+      "content": "<extra_id_190>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256110": {
+      "content": "<extra_id_189>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256111": {
+      "content": "<extra_id_188>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256112": {
+      "content": "<extra_id_187>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256113": {
+      "content": "<extra_id_186>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256114": {
+      "content": "<extra_id_185>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256115": {
+      "content": "<extra_id_184>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256116": {
+      "content": "<extra_id_183>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256117": {
+      "content": "<extra_id_182>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256118": {
+      "content": "<extra_id_181>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256119": {
+      "content": "<extra_id_180>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256120": {
+      "content": "<extra_id_179>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256121": {
+      "content": "<extra_id_178>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256122": {
+      "content": "<extra_id_177>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256123": {
+      "content": "<extra_id_176>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256124": {
+      "content": "<extra_id_175>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256125": {
+      "content": "<extra_id_174>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256126": {
+      "content": "<extra_id_173>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256127": {
+      "content": "<extra_id_172>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256128": {
+      "content": "<extra_id_171>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256129": {
+      "content": "<extra_id_170>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256130": {
+      "content": "<extra_id_169>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256131": {
+      "content": "<extra_id_168>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256132": {
+      "content": "<extra_id_167>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256133": {
+      "content": "<extra_id_166>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256134": {
+      "content": "<extra_id_165>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256135": {
+      "content": "<extra_id_164>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256136": {
+      "content": "<extra_id_163>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256137": {
+      "content": "<extra_id_162>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256138": {
+      "content": "<extra_id_161>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256139": {
+      "content": "<extra_id_160>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256140": {
+      "content": "<extra_id_159>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256141": {
+      "content": "<extra_id_158>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256142": {
+      "content": "<extra_id_157>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256143": {
+      "content": "<extra_id_156>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256144": {
+      "content": "<extra_id_155>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256145": {
+      "content": "<extra_id_154>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256146": {
+      "content": "<extra_id_153>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256147": {
+      "content": "<extra_id_152>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256148": {
+      "content": "<extra_id_151>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256149": {
+      "content": "<extra_id_150>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256150": {
+      "content": "<extra_id_149>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256151": {
+      "content": "<extra_id_148>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256152": {
+      "content": "<extra_id_147>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256153": {
+      "content": "<extra_id_146>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256154": {
+      "content": "<extra_id_145>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256155": {
+      "content": "<extra_id_144>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256156": {
+      "content": "<extra_id_143>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256157": {
+      "content": "<extra_id_142>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256158": {
+      "content": "<extra_id_141>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256159": {
+      "content": "<extra_id_140>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256160": {
+      "content": "<extra_id_139>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256161": {
+      "content": "<extra_id_138>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256162": {
+      "content": "<extra_id_137>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256163": {
+      "content": "<extra_id_136>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256164": {
+      "content": "<extra_id_135>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256165": {
+      "content": "<extra_id_134>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256166": {
+      "content": "<extra_id_133>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256167": {
+      "content": "<extra_id_132>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256168": {
+      "content": "<extra_id_131>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256169": {
+      "content": "<extra_id_130>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256170": {
+      "content": "<extra_id_129>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256171": {
+      "content": "<extra_id_128>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256172": {
+      "content": "<extra_id_127>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256173": {
+      "content": "<extra_id_126>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256174": {
+      "content": "<extra_id_125>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256175": {
+      "content": "<extra_id_124>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256176": {
+      "content": "<extra_id_123>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256177": {
+      "content": "<extra_id_122>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256178": {
+      "content": "<extra_id_121>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256179": {
+      "content": "<extra_id_120>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256180": {
+      "content": "<extra_id_119>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256181": {
+      "content": "<extra_id_118>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256182": {
+      "content": "<extra_id_117>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256183": {
+      "content": "<extra_id_116>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256184": {
+      "content": "<extra_id_115>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256185": {
+      "content": "<extra_id_114>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256186": {
+      "content": "<extra_id_113>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256187": {
+      "content": "<extra_id_112>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256188": {
+      "content": "<extra_id_111>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256189": {
+      "content": "<extra_id_110>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256190": {
+      "content": "<extra_id_109>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256191": {
+      "content": "<extra_id_108>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256192": {
+      "content": "<extra_id_107>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256193": {
+      "content": "<extra_id_106>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256194": {
+      "content": "<extra_id_105>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256195": {
+      "content": "<extra_id_104>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256196": {
+      "content": "<extra_id_103>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256197": {
+      "content": "<extra_id_102>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256198": {
+      "content": "<extra_id_101>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256199": {
+      "content": "<extra_id_100>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256200": {
+      "content": "<extra_id_99>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256201": {
+      "content": "<extra_id_98>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256202": {
+      "content": "<extra_id_97>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256203": {
+      "content": "<extra_id_96>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256204": {
+      "content": "<extra_id_95>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256205": {
+      "content": "<extra_id_94>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256206": {
+      "content": "<extra_id_93>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256207": {
+      "content": "<extra_id_92>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256208": {
+      "content": "<extra_id_91>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256209": {
+      "content": "<extra_id_90>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256210": {
+      "content": "<extra_id_89>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256211": {
+      "content": "<extra_id_88>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256212": {
+      "content": "<extra_id_87>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256213": {
+      "content": "<extra_id_86>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256214": {
+      "content": "<extra_id_85>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256215": {
+      "content": "<extra_id_84>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256216": {
+      "content": "<extra_id_83>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256217": {
+      "content": "<extra_id_82>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256218": {
+      "content": "<extra_id_81>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256219": {
+      "content": "<extra_id_80>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256220": {
+      "content": "<extra_id_79>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256221": {
+      "content": "<extra_id_78>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256222": {
+      "content": "<extra_id_77>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256223": {
+      "content": "<extra_id_76>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256224": {
+      "content": "<extra_id_75>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256225": {
+      "content": "<extra_id_74>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256226": {
+      "content": "<extra_id_73>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256227": {
+      "content": "<extra_id_72>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256228": {
+      "content": "<extra_id_71>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256229": {
+      "content": "<extra_id_70>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256230": {
+      "content": "<extra_id_69>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256231": {
+      "content": "<extra_id_68>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256232": {
+      "content": "<extra_id_67>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256233": {
+      "content": "<extra_id_66>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256234": {
+      "content": "<extra_id_65>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256235": {
+      "content": "<extra_id_64>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256236": {
+      "content": "<extra_id_63>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256237": {
+      "content": "<extra_id_62>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256238": {
+      "content": "<extra_id_61>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256239": {
+      "content": "<extra_id_60>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256240": {
+      "content": "<extra_id_59>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256241": {
+      "content": "<extra_id_58>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256242": {
+      "content": "<extra_id_57>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256243": {
+      "content": "<extra_id_56>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256244": {
+      "content": "<extra_id_55>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256245": {
+      "content": "<extra_id_54>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256246": {
+      "content": "<extra_id_53>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256247": {
+      "content": "<extra_id_52>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256248": {
+      "content": "<extra_id_51>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256249": {
+      "content": "<extra_id_50>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256250": {
+      "content": "<extra_id_49>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256251": {
+      "content": "<extra_id_48>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256252": {
+      "content": "<extra_id_47>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256253": {
+      "content": "<extra_id_46>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256254": {
+      "content": "<extra_id_45>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256255": {
+      "content": "<extra_id_44>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256256": {
+      "content": "<extra_id_43>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256257": {
+      "content": "<extra_id_42>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256258": {
+      "content": "<extra_id_41>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256259": {
+      "content": "<extra_id_40>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256260": {
+      "content": "<extra_id_39>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256261": {
+      "content": "<extra_id_38>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256262": {
+      "content": "<extra_id_37>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256263": {
+      "content": "<extra_id_36>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256264": {
+      "content": "<extra_id_35>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256265": {
+      "content": "<extra_id_34>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256266": {
+      "content": "<extra_id_33>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256267": {
+      "content": "<extra_id_32>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256268": {
+      "content": "<extra_id_31>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256269": {
+      "content": "<extra_id_30>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256270": {
+      "content": "<extra_id_29>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256271": {
+      "content": "<extra_id_28>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256272": {
+      "content": "<extra_id_27>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256273": {
+      "content": "<extra_id_26>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256274": {
+      "content": "<extra_id_25>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256275": {
+      "content": "<extra_id_24>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256276": {
+      "content": "<extra_id_23>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256277": {
+      "content": "<extra_id_22>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256278": {
+      "content": "<extra_id_21>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256279": {
+      "content": "<extra_id_20>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256280": {
+      "content": "<extra_id_19>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256281": {
+      "content": "<extra_id_18>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256282": {
+      "content": "<extra_id_17>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256283": {
+      "content": "<extra_id_16>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256284": {
+      "content": "<extra_id_15>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256285": {
+      "content": "<extra_id_14>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256286": {
+      "content": "<extra_id_13>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256287": {
+      "content": "<extra_id_12>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256288": {
+      "content": "<extra_id_11>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256289": {
+      "content": "<extra_id_10>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256290": {
+      "content": "<extra_id_9>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256291": {
+      "content": "<extra_id_8>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256292": {
+      "content": "<extra_id_7>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256293": {
+      "content": "<extra_id_6>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256294": {
+      "content": "<extra_id_5>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256295": {
+      "content": "<extra_id_4>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256296": {
+      "content": "<extra_id_3>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256297": {
+      "content": "<extra_id_2>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256298": {
+      "content": "<extra_id_1>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "256299": {
+      "content": "<extra_id_0>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "additional_special_tokens": [
+    "<extra_id_0>",
+    "<extra_id_1>",
+    "<extra_id_2>",
+    "<extra_id_3>",
+    "<extra_id_4>",
+    "<extra_id_5>",
+    "<extra_id_6>",
+    "<extra_id_7>",
+    "<extra_id_8>",
+    "<extra_id_9>",
+    "<extra_id_10>",
+    "<extra_id_11>",
+    "<extra_id_12>",
+    "<extra_id_13>",
+    "<extra_id_14>",
+    "<extra_id_15>",
+    "<extra_id_16>",
+    "<extra_id_17>",
+    "<extra_id_18>",
+    "<extra_id_19>",
+    "<extra_id_20>",
+    "<extra_id_21>",
+    "<extra_id_22>",
+    "<extra_id_23>",
+    "<extra_id_24>",
+    "<extra_id_25>",
+    "<extra_id_26>",
+    "<extra_id_27>",
+    "<extra_id_28>",
+    "<extra_id_29>",
+    "<extra_id_30>",
+    "<extra_id_31>",
+    "<extra_id_32>",
+    "<extra_id_33>",
+    "<extra_id_34>",
+    "<extra_id_35>",
+    "<extra_id_36>",
+    "<extra_id_37>",
+    "<extra_id_38>",
+    "<extra_id_39>",
+    "<extra_id_40>",
+    "<extra_id_41>",
+    "<extra_id_42>",
+    "<extra_id_43>",
+    "<extra_id_44>",
+    "<extra_id_45>",
+    "<extra_id_46>",
+    "<extra_id_47>",
+    "<extra_id_48>",
+    "<extra_id_49>",
+    "<extra_id_50>",
+    "<extra_id_51>",
+    "<extra_id_52>",
+    "<extra_id_53>",
+    "<extra_id_54>",
+    "<extra_id_55>",
+    "<extra_id_56>",
+    "<extra_id_57>",
+    "<extra_id_58>",
+    "<extra_id_59>",
+    "<extra_id_60>",
+    "<extra_id_61>",
+    "<extra_id_62>",
+    "<extra_id_63>",
+    "<extra_id_64>",
+    "<extra_id_65>",
+    "<extra_id_66>",
+    "<extra_id_67>",
+    "<extra_id_68>",
+    "<extra_id_69>",
+    "<extra_id_70>",
+    "<extra_id_71>",
+    "<extra_id_72>",
+    "<extra_id_73>",
+    "<extra_id_74>",
+    "<extra_id_75>",
+    "<extra_id_76>",
+    "<extra_id_77>",
+    "<extra_id_78>",
+    "<extra_id_79>",
+    "<extra_id_80>",
+    "<extra_id_81>",
+    "<extra_id_82>",
+    "<extra_id_83>",
+    "<extra_id_84>",
+    "<extra_id_85>",
+    "<extra_id_86>",
+    "<extra_id_87>",
+    "<extra_id_88>",
+    "<extra_id_89>",
+    "<extra_id_90>",
+    "<extra_id_91>",
+    "<extra_id_92>",
+    "<extra_id_93>",
+    "<extra_id_94>",
+    "<extra_id_95>",
+    "<extra_id_96>",
+    "<extra_id_97>",
+    "<extra_id_98>",
+    "<extra_id_99>",
+    "<extra_id_100>",
+    "<extra_id_101>",
+    "<extra_id_102>",
+    "<extra_id_103>",
+    "<extra_id_104>",
+    "<extra_id_105>",
+    "<extra_id_106>",
+    "<extra_id_107>",
+    "<extra_id_108>",
+    "<extra_id_109>",
+    "<extra_id_110>",
+    "<extra_id_111>",
+    "<extra_id_112>",
+    "<extra_id_113>",
+    "<extra_id_114>",
+    "<extra_id_115>",
+    "<extra_id_116>",
+    "<extra_id_117>",
+    "<extra_id_118>",
+    "<extra_id_119>",
+    "<extra_id_120>",
+    "<extra_id_121>",
+    "<extra_id_122>",
+    "<extra_id_123>",
+    "<extra_id_124>",
+    "<extra_id_125>",
+    "<extra_id_126>",
+    "<extra_id_127>",
+    "<extra_id_128>",
+    "<extra_id_129>",
+    "<extra_id_130>",
+    "<extra_id_131>",
+    "<extra_id_132>",
+    "<extra_id_133>",
+    "<extra_id_134>",
+    "<extra_id_135>",
+    "<extra_id_136>",
+    "<extra_id_137>",
+    "<extra_id_138>",
+    "<extra_id_139>",
+    "<extra_id_140>",
+    "<extra_id_141>",
+    "<extra_id_142>",
+    "<extra_id_143>",
+    "<extra_id_144>",
+    "<extra_id_145>",
+    "<extra_id_146>",
+    "<extra_id_147>",
+    "<extra_id_148>",
+    "<extra_id_149>",
+    "<extra_id_150>",
+    "<extra_id_151>",
+    "<extra_id_152>",
+    "<extra_id_153>",
+    "<extra_id_154>",
+    "<extra_id_155>",
+    "<extra_id_156>",
+    "<extra_id_157>",
+    "<extra_id_158>",
+    "<extra_id_159>",
+    "<extra_id_160>",
+    "<extra_id_161>",
+    "<extra_id_162>",
+    "<extra_id_163>",
+    "<extra_id_164>",
+    "<extra_id_165>",
+    "<extra_id_166>",
+    "<extra_id_167>",
+    "<extra_id_168>",
+    "<extra_id_169>",
+    "<extra_id_170>",
+    "<extra_id_171>",
+    "<extra_id_172>",
+    "<extra_id_173>",
+    "<extra_id_174>",
+    "<extra_id_175>",
+    "<extra_id_176>",
+    "<extra_id_177>",
+    "<extra_id_178>",
+    "<extra_id_179>",
+    "<extra_id_180>",
+    "<extra_id_181>",
+    "<extra_id_182>",
+    "<extra_id_183>",
+    "<extra_id_184>",
+    "<extra_id_185>",
+    "<extra_id_186>",
+    "<extra_id_187>",
+    "<extra_id_188>",
+    "<extra_id_189>",
+    "<extra_id_190>",
+    "<extra_id_191>",
+    "<extra_id_192>",
+    "<extra_id_193>",
+    "<extra_id_194>",
+    "<extra_id_195>",
+    "<extra_id_196>",
+    "<extra_id_197>",
+    "<extra_id_198>",
+    "<extra_id_199>",
+    "<extra_id_200>",
+    "<extra_id_201>",
+    "<extra_id_202>",
+    "<extra_id_203>",
+    "<extra_id_204>",
+    "<extra_id_205>",
+    "<extra_id_206>",
+    "<extra_id_207>",
+    "<extra_id_208>",
+    "<extra_id_209>",
+    "<extra_id_210>",
+    "<extra_id_211>",
+    "<extra_id_212>",
+    "<extra_id_213>",
+    "<extra_id_214>",
+    "<extra_id_215>",
+    "<extra_id_216>",
+    "<extra_id_217>",
+    "<extra_id_218>",
+    "<extra_id_219>",
+    "<extra_id_220>",
+    "<extra_id_221>",
+    "<extra_id_222>",
+    "<extra_id_223>",
+    "<extra_id_224>",
+    "<extra_id_225>",
+    "<extra_id_226>",
+    "<extra_id_227>",
+    "<extra_id_228>",
+    "<extra_id_229>",
+    "<extra_id_230>",
+    "<extra_id_231>",
+    "<extra_id_232>",
+    "<extra_id_233>",
+    "<extra_id_234>",
+    "<extra_id_235>",
+    "<extra_id_236>",
+    "<extra_id_237>",
+    "<extra_id_238>",
+    "<extra_id_239>",
+    "<extra_id_240>",
+    "<extra_id_241>",
+    "<extra_id_242>",
+    "<extra_id_243>",
+    "<extra_id_244>",
+    "<extra_id_245>",
+    "<extra_id_246>",
+    "<extra_id_247>",
+    "<extra_id_248>",
+    "<extra_id_249>",
+    "<extra_id_250>",
+    "<extra_id_251>",
+    "<extra_id_252>",
+    "<extra_id_253>",
+    "<extra_id_254>",
+    "<extra_id_255>",
+    "<extra_id_256>",
+    "<extra_id_257>",
+    "<extra_id_258>",
+    "<extra_id_259>",
+    "<extra_id_260>",
+    "<extra_id_261>",
+    "<extra_id_262>",
+    "<extra_id_263>",
+    "<extra_id_264>",
+    "<extra_id_265>",
+    "<extra_id_266>",
+    "<extra_id_267>",
+    "<extra_id_268>",
+    "<extra_id_269>",
+    "<extra_id_270>",
+    "<extra_id_271>",
+    "<extra_id_272>",
+    "<extra_id_273>",
+    "<extra_id_274>",
+    "<extra_id_275>",
+    "<extra_id_276>",
+    "<extra_id_277>",
+    "<extra_id_278>",
+    "<extra_id_279>",
+    "<extra_id_280>",
+    "<extra_id_281>",
+    "<extra_id_282>",
+    "<extra_id_283>",
+    "<extra_id_284>",
+    "<extra_id_285>",
+    "<extra_id_286>",
+    "<extra_id_287>",
+    "<extra_id_288>",
+    "<extra_id_289>",
+    "<extra_id_290>",
+    "<extra_id_291>",
+    "<extra_id_292>",
+    "<extra_id_293>",
+    "<extra_id_294>",
+    "<extra_id_295>",
+    "<extra_id_296>",
+    "<extra_id_297>",
+    "<extra_id_298>",
+    "<extra_id_299>"
+  ],
+  "bos_token": "<s>",
+  "clean_up_tokenization_spaces": true,
+  "eos_token": "</s>",
+  "extra_ids": 300,
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "<pad>",
+  "sp_model_kwargs": {},
+  "spaces_between_special_tokens": false,
+  "tokenizer_class": "T5Tokenizer",
+  "unk_token": "<unk>"
+}