{ "cells": [ { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import json\n", "import os\n", "\n", "# os.environ[\"TORCH_LOGS\"] = \"+dynamo,output_code,graph_breaks,recompiles\"\n", "\n", "import torch\n", "from diffusers import CogVideoXPipeline, CogVideoXDDIMScheduler\n", "from diffusers.utils import export_to_video\n", "from torchao.quantization import (\n", " autoquant,\n", " quantize_,\n", " int8_weight_only,\n", " int8_dynamic_activation_int8_weight,\n", " int8_dynamic_activation_int4_weight,\n", " int8_dynamic_activation_int8_semi_sparse_weight,\n", " int4_weight_only,\n", " float8_dynamic_activation_float8_weight,\n", " float8_weight_only,\n", " fpx_weight_only,\n", ")\n", "from torchao.quantization.quant_api import PerRow\n", "from torchao.sparsity import sparsify_\n", "\n", "from diffusers_torchao.inference.utils import (\n", " cleanup_tmp_directory,\n", " benchmark_fn,\n", " pretty_print_results,\n", " print_memory,\n", " reset_memory,\n", ")\n", "\n", "# Set high precision for float32 matrix multiplications.\n", "# This setting optimizes performance on NVIDIA GPUs with Ampere architecture (e.g., A100, RTX 30 series) or newer.\n", "torch.set_float32_matmul_precision(\"high\")\n", "\n", "\n", "CONVERT_DTYPE = {\n", " \"fp16\": lambda module: module.to(dtype=torch.float16),\n", " \"bf16\": lambda module: module.to(dtype=torch.bfloat16),\n", " \"fp8wo\": lambda module: quantize_(module, float8_weight_only()),\n", " \"fp8dq\": lambda module: quantize_(\n", " module, float8_dynamic_activation_float8_weight()\n", " ),\n", " \"fp8dqrow\": lambda module: quantize_(\n", " module, float8_dynamic_activation_float8_weight(granularity=PerRow())\n", " ),\n", " \"fp6_e3m2\": lambda module: quantize_(module, fpx_weight_only(3, 2)),\n", " \"fp5_e2m2\": lambda module: quantize_(module, fpx_weight_only(2, 2)),\n", " \"fp4_e2m1\": lambda module: quantize_(module, fpx_weight_only(2, 1)),\n", " \"int8wo\": lambda module: quantize_(module, int8_weight_only()),\n", " \"int8dq\": lambda module: quantize_(module, int8_dynamic_activation_int8_weight()),\n", " \"int4dq\": lambda module: quantize_(module, int8_dynamic_activation_int4_weight()),\n", " \"int4wo\": lambda module: quantize_(module, int4_weight_only()),\n", " \"autoquant\": lambda module: autoquant(module, error_on_unseen=False),\n", " \"sparsify\": lambda module: sparsify_(\n", " module, int8_dynamic_activation_int8_semi_sparse_weight()\n", " ),\n", "}\n", "\n", "\n", "def load_pipeline(model_id, dtype, device, quantize_vae, compile, fuse_qkv):\n", " # 1. Load pipeline\n", " pipe = CogVideoXPipeline.from_pretrained(model_id, torch_dtype=torch.bfloat16).to(\n", " device\n", " )\n", " pipe.scheduler = CogVideoXDDIMScheduler.from_config(\n", " pipe.scheduler.config, timestep_spacing=\"trailing\"\n", " )\n", " pipe.set_progress_bar_config(disable=True)\n", "\n", " if fuse_qkv:\n", " pipe.fuse_qkv_projections()\n", "\n", " # 2. Quantize and compile\n", " if dtype == \"autoquant\" and compile:\n", " pipe.transformer.to(memory_format=torch.channels_last)\n", " pipe.transformer = torch.compile(\n", " pipe.transformer, mode=\"max-autotune\", fullgraph=True\n", " )\n", " # VAE cannot be compiled due to: https://gist.github.com/a-r-r-o-w/5183d75e452a368fd17448fcc810bd3f#file-test_cogvideox_torch_compile-py-L30\n", "\n", " text_encoder_return = CONVERT_DTYPE[dtype](pipe.text_encoder)\n", " transformer_return = CONVERT_DTYPE[dtype](pipe.transformer)\n", " vae_return = None\n", " if dtype in [\"fp32\", \"fp16\", \"bf16\", \"fp8_e4m3\", \"fp8_e5m2\"] or quantize_vae:\n", " vae_return = CONVERT_DTYPE[dtype](pipe.vae)\n", "\n", " if text_encoder_return is not None:\n", " pipe.text_encoder = text_encoder_return\n", " if transformer_return is not None:\n", " pipe.transformer = transformer_return\n", " if vae_return is not None:\n", " pipe.vae = vae_return\n", "\n", " if dtype != \"autoquant\" and compile:\n", " pipe.transformer.to(memory_format=torch.channels_last)\n", " pipe.transformer = torch.compile(\n", " pipe.transformer, mode=\"max-autotune\", fullgraph=True\n", " )\n", " # VAE cannot be compiled due to: https://gist.github.com/a-r-r-o-w/5183d75e452a368fd17448fcc810bd3f#file-test_cogvideox_torch_compile-py-L30\n", "\n", " return pipe\n", "\n", "\n", "def run_inference(pipe):\n", " prompt = (\n", " \"A panda, dressed in a small, red jacket and a tiny hat, sits on a wooden stool in a serene bamboo forest. \"\n", " \"The panda's fluffy paws strum a miniature acoustic guitar, producing soft, melodic tunes. Nearby, a few other \"\n", " \"pandas gather, watching curiously and some clapping in rhythm. Sunlight filters through the tall bamboo, \"\n", " \"casting a gentle glow on the scene. The panda's face is expressive, showing concentration and joy as it plays. \"\n", " \"The background includes a small, flowing stream and vibrant green foliage, enhancing the peaceful and magical \"\n", " \"atmosphere of this unique musical performance.\"\n", " )\n", " guidance_scale = 6\n", " num_inference_steps = 50\n", "\n", " video = pipe(\n", " prompt=prompt,\n", " guidance_scale=guidance_scale,\n", " use_dynamic_cfg=True,\n", " num_inference_steps=num_inference_steps,\n", " generator=torch.Generator().manual_seed(\n", " 3047\n", " ), # https://arxiv.org/abs/2109.08203\n", " )\n", " return video\n", "\n", "\n", "def main(model_id, dtype, device, quantize_vae, compile, fuse_qkv):\n", " reset_memory(device)\n", "\n", " # 1. Load pipeline\n", " pipe = load_pipeline(model_id, dtype, device, quantize_vae, compile, fuse_qkv)\n", "\n", " print_memory(device)\n", "\n", " torch.cuda.empty_cache()\n", " model_memory = round(torch.cuda.memory_allocated() / 1024**3, 3)\n", "\n", " # 2. Warmup\n", " num_warmups = 2\n", " for _ in range(num_warmups):\n", " video = run_inference(pipe)\n", "\n", " # 3. Benchmark\n", " time = benchmark_fn(run_inference, pipe)\n", " print_memory(device)\n", "\n", " torch.cuda.empty_cache()\n", " inference_memory = round(torch.cuda.max_memory_allocated() / 1024**3, 3)\n", "\n", " # 4. Save results\n", " model_type = \"5B\" if \"5b\" in model_id else \"2B\"\n", " info = {\n", " \"model_type\": model_type,\n", " \"compile\": compile,\n", " \"fuse_qkv\": fuse_qkv,\n", " \"quantize_vae\": quantize_vae,\n", " \"quantization\": dtype,\n", " \"model_memory\": model_memory,\n", " \"inference_memory\": inference_memory,\n", " \"time\": time,\n", " }\n", " pretty_print_results(info, precision=3)\n", "\n", " # Serialize artifacts\n", " model_name = model_id.replace(\"/\", \"_\").replace(\".\", \"_\")\n", " filename_prefix = f\"output-model_{model_name}-quantization_{dtype}-compile_{compile}-fuse_qkv_{fuse_qkv}\"\n", "\n", " with open(f\"{filename_prefix}.json\", \"w\") as file:\n", " json.dump(info, file)\n", "\n", " export_to_video(video.frames[0], f\"{filename_prefix}.mp4\", fps=8)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "model_id = \"THUDM/CogVideoX-5b\"\n", "dtype = \"fp8dqrow\"\n", "device = \"cuda:0\"\n", "quantize_vae = False\n", "compile = True\n", "fuse_qkv = False\n", "\n", "main(\n", " model_id=model_id,\n", " dtype=dtype,\n", " device=device,\n", " quantize_vae=quantize_vae,\n", " compile=compile,\n", " fuse_qkv=fuse_qkv,\n", ")\n", "cleanup_tmp_directory()" ] } ], "metadata": { "kernelspec": { "display_name": "suno_env", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.10.14" } }, "nbformat": 4, "nbformat_minor": 2 }