Upload root source files

.gitignore

@@ -0,0 +1,5 @@
+__pycache__
+*.egg-info
+
+wan_models
+checkpoints

LICENSE

@@ -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.

README.md

@@ -0,0 +1,109 @@
+<p align="center">
+<h1 align="center">Self Forcing</h1>
+<h3 align="center">Bridging the Train-Test Gap in Autoregressive Video Diffusion</h3>
+</p>
+<p align="center">
+  <p align="center">
+    <a href="https://www.xunhuang.me/">Xun Huang</a><sup>1</sup>
+    ·
+    <a href="https://zhengqili.github.io/">Zhengqi Li</a><sup>1</sup>
+    ·
+    <a href="https://guandehe.github.io/">Guande He</a><sup>2</sup>
+    ·
+    <a href="https://mingyuanzhou.github.io/">Mingyuan Zhou</a><sup>2</sup>
+    ·
+    <a href="https://research.adobe.com/person/eli-shechtman/">Eli Shechtman</a><sup>1</sup><br>
+    <sup>1</sup>Adobe Research <sup>2</sup>UT Austin
+  </p>
+  <h3 align="center"><a href="https://arxiv.org/abs/2506.08009">Paper</a> | <a href="https://self-forcing.github.io">Website</a> | <a href="https://huggingface.co/gdhe17/Self-Forcing/tree/main">Models (HuggingFace)</a></h3>
+</p>
+
+---
+
+Self Forcing trains autoregressive video diffusion models by **simulating the inference process during training**, performing autoregressive rollout with KV caching. It resolves the train-test distribution mismatch and enables **real-time, streaming video generation on a single RTX 4090** while matching the quality of state-of-the-art diffusion models.
+
+---
+
+
+https://github.com/user-attachments/assets/7548c2db-fe03-4ba8-8dd3-52d2c6160739
+
+
+## Requirements
+We tested this repo on the following setup:
+* Nvidia GPU with at least 24 GB memory (RTX 4090, A100, and H100 are tested).
+* Linux operating system.
+* 64 GB RAM.
+
+Other hardware setup could also work but hasn't been tested.
+
+## Installation
+Create a conda environment and install dependencies:
+```
+conda create -n self_forcing python=3.10 -y
+conda activate self_forcing
+pip install -r requirements.txt
+pip install flash-attn --no-build-isolation
+python setup.py develop
+```
+
+## Quick Start
+### Download checkpoints
+```
+huggingface-cli download Wan-AI/Wan2.1-T2V-1.3B --local-dir-use-symlinks False --local-dir wan_models/Wan2.1-T2V-1.3B
+huggingface-cli download gdhe17/Self-Forcing checkpoints/self_forcing_dmd.pt --local-dir .
+```
+
+### GUI demo
+```
+python demo.py
+```
+Note:
+* **Our model works better with long, detailed prompts** since it's trained with such prompts. We will integrate prompt extension into the codebase (similar to [Wan2.1](https://github.com/Wan-Video/Wan2.1/tree/main?tab=readme-ov-file#2-using-prompt-extention)) in the future. For now, it is recommended to use third-party LLMs (such as GPT-4o) to extend your prompt before providing to the model.
+* You may want to adjust FPS so it plays smoothly on your device.
+* The speed can be improved by enabling `torch.compile`, [TAEHV-VAE](https://github.com/madebyollin/taehv/), or using FP8 Linear layers, although the latter two options may sacrifice quality. It is recommended to use `torch.compile` if possible and enable TAEHV-VAE if further speedup is needed.
+
+### CLI Inference
+Example inference script using the chunk-wise autoregressive checkpoint trained with DMD:
+```
+python inference.py \
+    --config_path configs/self_forcing_dmd.yaml \
+    --output_folder videos/self_forcing_dmd \
+    --checkpoint_path checkpoints/self_forcing_dmd.pt \
+    --data_path prompts/MovieGenVideoBench_extended.txt \
+    --use_ema
+```
+Other config files and corresponding checkpoints can be found in [configs](configs) folder and our [huggingface repo](https://huggingface.co/gdhe17/Self-Forcing/tree/main/checkpoints).
+
+## Training
+### Download text prompts and ODE initialized checkpoint
+```
+huggingface-cli download gdhe17/Self-Forcing checkpoints/ode_init.pt --local-dir .
+huggingface-cli download gdhe17/Self-Forcing vidprom_filtered_extended.txt --local-dir prompts
+```
+Note: Our training algorithm (except for the GAN version) is data-free (**no video data is needed**). For now, we directly provide the ODE initialization checkpoint and will add more instructions on how to perform ODE initialization in the future (which is identical to the process described in the [CausVid](https://github.com/tianweiy/CausVid) repo).
+
+### Self Forcing Training with DMD
+```
+torchrun --nnodes=8 --nproc_per_node=8 --rdzv_id=5235 \
+  --rdzv_backend=c10d \
+  --rdzv_endpoint $MASTER_ADDR \
+  train.py \
+  --config_path configs/self_forcing_dmd.yaml \
+  --logdir logs/self_forcing_dmd \
+  --disable-wandb
+```
+Our training run uses 600 iterations and completes in under 2 hours using 64 H100 GPUs. By implementing gradient accumulation, it should be possible to reproduce the results in less than 16 hours using 8 H100 GPUs.
+
+## Acknowledgements
+This codebase is built on top of the open-source implementation of [CausVid](https://github.com/tianweiy/CausVid) by [Tianwei Yin](https://tianweiy.github.io/) and the [Wan2.1](https://github.com/Wan-Video/Wan2.1) repo.
+
+## Citation
+If you find this codebase useful for your research, please kindly cite our paper:
+```
+@article{huang2025selfforcing,
+  title={Self Forcing: Bridging the Train-Test Gap in Autoregressive Video Diffusion},
+  author={Huang, Xun and Li, Zhengqi and He, Guande and Zhou, Mingyuan and Shechtman, Eli},
+  journal={arXiv preprint arXiv:2506.08009},
+  year={2025}
+}
+```

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)

inference.py

@@ -0,0 +1,192 @@
+import argparse
+import torch
+import os
+from omegaconf import OmegaConf
+from tqdm import tqdm
+from torchvision import transforms
+from torchvision.io import write_video
+from einops import rearrange
+import torch.distributed as dist
+from torch.utils.data import DataLoader, SequentialSampler
+from torch.utils.data.distributed import DistributedSampler
+
+from pipeline import (
+    CausalDiffusionInferencePipeline,
+    CausalInferencePipeline,
+)
+from utils.dataset import TextDataset, TextImagePairDataset
+from utils.misc import set_seed
+
+from demo_utils.memory import gpu, get_cuda_free_memory_gb, DynamicSwapInstaller
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--config_path", type=str, help="Path to the config file")
+parser.add_argument("--checkpoint_path", type=str, help="Path to the checkpoint folder")
+parser.add_argument("--data_path", type=str, help="Path to the dataset")
+parser.add_argument("--extended_prompt_path", type=str, help="Path to the extended prompt")
+parser.add_argument("--output_folder", type=str, help="Output folder")
+parser.add_argument("--num_output_frames", type=int, default=21,
+                    help="Number of overlap frames between sliding windows")
+parser.add_argument("--i2v", action="store_true", help="Whether to perform I2V (or T2V by default)")
+parser.add_argument("--use_ema", action="store_true", help="Whether to use EMA parameters")
+parser.add_argument("--seed", type=int, default=0, help="Random seed")
+parser.add_argument("--num_samples", type=int, default=1, help="Number of samples to generate per prompt")
+parser.add_argument("--save_with_index", action="store_true",
+                    help="Whether to save the video using the index or prompt as the filename")
+args = parser.parse_args()
+
+# Initialize distributed inference
+if "LOCAL_RANK" in os.environ:
+    dist.init_process_group(backend='nccl')
+    local_rank = int(os.environ["LOCAL_RANK"])
+    torch.cuda.set_device(local_rank)
+    device = torch.device(f"cuda:{local_rank}")
+    world_size = dist.get_world_size()
+    set_seed(args.seed + local_rank)
+else:
+    device = torch.device("cuda")
+    local_rank = 0
+    world_size = 1
+    set_seed(args.seed)
+
+print(f'Free VRAM {get_cuda_free_memory_gb(gpu)} GB')
+low_memory = get_cuda_free_memory_gb(gpu) < 40
+
+torch.set_grad_enabled(False)
+
+config = OmegaConf.load(args.config_path)
+default_config = OmegaConf.load("configs/default_config.yaml")
+config = OmegaConf.merge(default_config, config)
+
+# Initialize pipeline
+if hasattr(config, 'denoising_step_list'):
+    # Few-step inference
+    pipeline = CausalInferencePipeline(config, device=device)
+else:
+    # Multi-step diffusion inference
+    pipeline = CausalDiffusionInferencePipeline(config, device=device)
+
+if args.checkpoint_path:
+    state_dict = torch.load(args.checkpoint_path, map_location="cpu")
+    pipeline.generator.load_state_dict(state_dict['generator' if not args.use_ema else 'generator_ema'])
+
+pipeline = pipeline.to(dtype=torch.bfloat16)
+if low_memory:
+    DynamicSwapInstaller.install_model(pipeline.text_encoder, device=gpu)
+else:
+    pipeline.text_encoder.to(device=gpu)
+pipeline.generator.to(device=gpu)
+pipeline.vae.to(device=gpu)
+
+
+# Create dataset
+if args.i2v:
+    assert not dist.is_initialized(), "I2V does not support distributed inference yet"
+    transform = transforms.Compose([
+        transforms.Resize((480, 832)),
+        transforms.ToTensor(),
+        transforms.Normalize([0.5], [0.5])
+    ])
+    dataset = TextImagePairDataset(args.data_path, transform=transform)
+else:
+    dataset = TextDataset(prompt_path=args.data_path, extended_prompt_path=args.extended_prompt_path)
+num_prompts = len(dataset)
+print(f"Number of prompts: {num_prompts}")
+
+if dist.is_initialized():
+    sampler = DistributedSampler(dataset, shuffle=False, drop_last=True)
+else:
+    sampler = SequentialSampler(dataset)
+dataloader = DataLoader(dataset, batch_size=1, sampler=sampler, num_workers=0, drop_last=False)
+
+# Create output directory (only on main process to avoid race conditions)
+if local_rank == 0:
+    os.makedirs(args.output_folder, exist_ok=True)
+
+if dist.is_initialized():
+    dist.barrier()
+
+
+def encode(self, videos: torch.Tensor) -> torch.Tensor:
+    device, dtype = videos[0].device, videos[0].dtype
+    scale = [self.mean.to(device=device, dtype=dtype),
+             1.0 / self.std.to(device=device, dtype=dtype)]
+    output = [
+        self.model.encode(u.unsqueeze(0), scale).float().squeeze(0)
+        for u in videos
+    ]
+
+    output = torch.stack(output, dim=0)
+    return output
+
+
+for i, batch_data in tqdm(enumerate(dataloader), disable=(local_rank != 0)):
+    idx = batch_data['idx'].item()
+
+    # For DataLoader batch_size=1, the batch_data is already a single item, but in a batch container
+    # Unpack the batch data for convenience
+    if isinstance(batch_data, dict):
+        batch = batch_data
+    elif isinstance(batch_data, list):
+        batch = batch_data[0]  # First (and only) item in the batch
+
+    all_video = []
+    num_generated_frames = 0  # Number of generated (latent) frames
+
+    if args.i2v:
+        # For image-to-video, batch contains image and caption
+        prompt = batch['prompts'][0]  # Get caption from batch
+        prompts = [prompt] * args.num_samples
+
+        # Process the image
+        image = batch['image'].squeeze(0).unsqueeze(0).unsqueeze(2).to(device=device, dtype=torch.bfloat16)
+
+        # Encode the input image as the first latent
+        initial_latent = pipeline.vae.encode_to_latent(image).to(device=device, dtype=torch.bfloat16)
+        initial_latent = initial_latent.repeat(args.num_samples, 1, 1, 1, 1)
+
+        sampled_noise = torch.randn(
+            [args.num_samples, args.num_output_frames - 1, 16, 60, 104], device=device, dtype=torch.bfloat16
+        )
+    else:
+        # For text-to-video, batch is just the text prompt
+        prompt = batch['prompts'][0]
+        extended_prompt = batch['extended_prompts'][0] if 'extended_prompts' in batch else None
+        if extended_prompt is not None:
+            prompts = [extended_prompt] * args.num_samples
+        else:
+            prompts = [prompt] * args.num_samples
+        initial_latent = None
+
+        sampled_noise = torch.randn(
+            [args.num_samples, args.num_output_frames, 16, 60, 104], device=device, dtype=torch.bfloat16
+        )
+
+    # Generate 81 frames
+    video, latents = pipeline.inference(
+        noise=sampled_noise,
+        text_prompts=prompts,
+        return_latents=True,
+        initial_latent=initial_latent,
+        low_memory=low_memory,
+    )
+    current_video = rearrange(video, 'b t c h w -> b t h w c').cpu()
+    all_video.append(current_video)
+    num_generated_frames += latents.shape[1]
+
+    # Final output video
+    video = 255.0 * torch.cat(all_video, dim=1)
+
+    # Clear VAE cache
+    pipeline.vae.model.clear_cache()
+
+    # Save the video if the current prompt is not a dummy prompt
+    if idx < num_prompts:
+        model = "regular" if not args.use_ema else "ema"
+        for seed_idx in range(args.num_samples):
+            # All processes save their videos
+            if args.save_with_index:
+                output_path = os.path.join(args.output_folder, f'{idx}-{seed_idx}_{model}.mp4')
+            else:
+                output_path = os.path.join(args.output_folder, f'{prompt[:100]}-{seed_idx}.mp4')
+            write_video(output_path, video[seed_idx], fps=16)

requirements.txt

@@ -0,0 +1,40 @@
+torch>=2.4.0
+torchvision>=0.19.0
+opencv-python>=4.9.0.80
+diffusers==0.31.0
+transformers>=4.49.0
+tokenizers>=0.20.3
+accelerate>=1.1.1
+tqdm
+imageio
+easydict
+ftfy
+dashscope
+imageio-ffmpeg
+numpy==1.24.4
+wandb
+omegaconf
+einops
+av==13.1.0
+opencv-python
+git+https://github.com/openai/CLIP.git
+open_clip_torch
+starlette
+pycocotools
+lmdb
+matplotlib
+sentencepiece
+pydantic==2.10.6
+scikit-image
+huggingface_hub[cli]
+dominate
+nvidia-pyindex
+nvidia-tensorrt
+pycuda
+onnx
+onnxruntime
+onnxscript
+onnxconverter_common
+flask
+flask-socketio
+torchao

setup.py

@@ -0,0 +1,6 @@
+from setuptools import setup, find_packages
+setup(
+    name="self_forcing",
+    version="0.0.1",
+    packages=find_packages(),
+)

train.py

@@ -0,0 +1,47 @@
+import argparse
+import os
+from omegaconf import OmegaConf
+import wandb
+
+from trainer import DiffusionTrainer, GANTrainer, ODETrainer, ScoreDistillationTrainer
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config_path", type=str, required=True)
+    parser.add_argument("--no_save", action="store_true")
+    parser.add_argument("--no_visualize", action="store_true")
+    parser.add_argument("--logdir", type=str, default="", help="Path to the directory to save logs")
+    parser.add_argument("--wandb-save-dir", type=str, default="", help="Path to the directory to save wandb logs")
+    parser.add_argument("--disable-wandb", action="store_true")
+
+    args = parser.parse_args()
+
+    config = OmegaConf.load(args.config_path)
+    default_config = OmegaConf.load("configs/default_config.yaml")
+    config = OmegaConf.merge(default_config, config)
+    config.no_save = args.no_save
+    config.no_visualize = args.no_visualize
+
+    # get the filename of config_path
+    config_name = os.path.basename(args.config_path).split(".")[0]
+    config.config_name = config_name
+    config.logdir = args.logdir
+    config.wandb_save_dir = args.wandb_save_dir
+    config.disable_wandb = args.disable_wandb
+
+    if config.trainer == "diffusion":
+        trainer = DiffusionTrainer(config)
+    elif config.trainer == "gan":
+        trainer = GANTrainer(config)
+    elif config.trainer == "ode":
+        trainer = ODETrainer(config)
+    elif config.trainer == "score_distillation":
+        trainer = ScoreDistillationTrainer(config)
+    trainer.train()
+
+    wandb.finish()
+
+
+if __name__ == "__main__":
+    main()