a-r-r-o-w

#VERBOSE=0 torchrun --nproc_per_node 3 self_contained_pp_LOC.py
import os, random, numpy as np, torch, torch.nn as nn, torch.distributed as dist, torch.nn.functional as F
from torch.optim import AdamW
from torch.utils.data import DataLoader, DistributedSampler
from datasets import load_dataset
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer

STEP, local_rank, world_size, verbose = 0, int(os.environ["LOCAL_RANK"]), int(os.environ["WORLD_SIZE"]), os.environ.get("VERBOSE", "0") == "1"

def set_all_seed(seed):

a-r-r-o-w

CogVideoX Official

• Text/Video-to-Video 2B: https://huggingface.co/THUDM/CogVideoX-2b
• Text/Video-to-Video 5B: https://huggingface.co/THUDM/CogVideoX-5b
• Image-to-Video 5B: https://huggingface.co/THUDM/CogVideoX-5b-I2V
• Original Repository: https://github.com/THUDM/CogVideo
• Diffusers documentation: https://huggingface.co/docs/diffusers/en/api/pipelines/cogvideox
• Diffusers-TorchAO quantization benchmarks: https://github.com/sayakpaul/diffusers-torchao/
• Diffusers-Quanto example: https://gist.github.com/a-r-r-o-w/31be62828b00a9292821b85c1017effa
• HF CogVideoX Space: https://huggingface.co/spaces/THUDM/CogVideoX-5B-Space

a-r-r-o-w

import gc

import torch
from accelerate.utils import compute_module_sizes
from diffusers import BitsAndBytesConfig, CogVideoXPipeline, CogVideoXTransformer3DModel
from diffusers.utils import export_to_video
from transformers import T5EncoderModel


def reset_memory():

a-r-r-o-w

#!/usr/bin/env python3

import argparse
import gc
import json
import os
import pathlib
import random
from typing import Any, Dict

a-r-r-o-w

import argparse
import os
import random
from typing import Any, Dict

import numpy as np
import pandas as pd
import torch
import transformers
from decord import VideoReader, cpu

a-r-r-o-w

# Install torchao from source and Pytorch Nightly
# Other environments have not yet been tested.

import gc

import torch
from diffusers import AutoencoderKLCogVideoX, CogVideoXTransformer3DModel, CogVideoXPipeline
from diffusers.utils import export_to_video
from transformers import T5EncoderModel
from torchao.quantization import (

a-r-r-o-w

#!/bin/bash

compile_flags=("" "--compile")
fuse_qkv_flags=("" "--fuse_qkv")
# quantizations=("fp16" "bf16" "fp8" "fp8_e4m3" "fp8_e5m2" "fp6" "int8wo" "int8dq" "int4dq" "int4wo" "autoquant" "sparsify")
quantizations=("fp16" "bf16" "fp6" "int8wo" "int8dq" "int4dq" "int4wo" "autoquant" "sparsify")
device="cuda"

# Check if completed.txt exists and read it into an array
if [ -f completed.txt ]; then

a-r-r-o-w

import gc
import torch
from diffusers import CogVideoXPipeline
from diffusers.utils import export_to_video
from optimum.quanto import freeze, quantize, qfloat8, qfloat8_e4m3fn, qfloat8_e5m2, qint8, qint4, qint2


def reset_memory(device):
    gc.collect()
    torch.cuda.empty_cache()

a-r-r-o-w

import argparse
import gc
import os
import time

os.environ["TORCH_LOGS"] = "dynamo"

import torch
from diffusers import AutoencoderKLCogVideoX, CogVideoXPipeline, CogVideoXTransformer3DModel, CogVideoXDDIMScheduler
from diffusers.utils import export_to_video

a-r-r-o-w

import gc

import torch
from diffusers import CogVideoXPipeline, CogVideoXDDIMScheduler
from diffusers.utils import export_to_video


def reset_memory():
    gc.collect()
    torch.cuda.empty_cache()