Maximofn commited on
Commit
11998d6
·
1 Parent(s): 086bafc

feat(src): :rocket: Add example script

Browse files
Files changed (1) hide show
  1. quantization.py +26 -0
quantization.py ADDED
@@ -0,0 +1,26 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
3
+ from diffusers import HunyuanVideoTransformer3DModel, HunyuanVideoPipeline
4
+ from diffusers.utils import export_to_video
5
+
6
+ quant_config = DiffusersBitsAndBytesConfig(load_in_8bit=True)
7
+ # model = "tencent/HunyuanVideo"
8
+ model = "hunyuanvideo-community/HunyuanVideo"
9
+ transformer_8bit = HunyuanVideoTransformer3DModel.from_pretrained(
10
+ model,
11
+ subfolder="transformer",
12
+ quantization_config=quant_config,
13
+ torch_dtype=torch.float16,
14
+ # device_map="balanced",
15
+ )
16
+
17
+ pipeline = HunyuanVideoPipeline.from_pretrained(
18
+ "hunyuanvideo-community/HunyuanVideo",
19
+ transformer=transformer_8bit,
20
+ torch_dtype=torch.float16,
21
+ device_map="balanced",
22
+ )
23
+
24
+ prompt = "A cat walks on the grass, realistic style."
25
+ video = pipeline(prompt=prompt, num_frames=61, num_inference_steps=3).frames[0]
26
+ export_to_video(video, "cat.mp4", fps=15)