nodes/kijai-ComfyUI-PyramidFlowWrapper
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Update pyramid_dit_for_video_gen_pipeline.py

Browse Source 
Change "torch.float8_e4m3fn, torch.float8_e4m3fn" to "torch.float8_e4m3fn, torch.float8_e5m2".
This commit is contained in:
gluttony-10
2024-10-11 10:25:45 +08:00
committed by GitHub
parent 7efb654cdb
commit dd58129e07
1 changed files with 3 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pyramid_dit/pyramid_dit_for_video_gen_pipeline.py
+3 -3
View File
@@ -47,7 +47,7 @@ class PyramidDiTForVideoGeneration:
):
super().__init__()
if model_dtype in [torch.float8_e4m3fn, torch.float8_e4m3fn]:
if model_dtype in [torch.float8_e4m3fn, torch.float8_e5m2]:
self.dtype = torch.bfloat16
else:
self.dtype = model_dtype
@@ -70,12 +70,12 @@ class PyramidDiTForVideoGeneration:
use_temporal_causal=True if not use_flash_attn else False,
interp_condition_pos=interp_condition_pos,
)
if model_dtype in [torch.float8_e4m3fn, torch.float8_e4m3fn]:
if model_dtype in [torch.float8_e4m3fn, torch.float8_e5m2]:
for name, param in self.dit.named_parameters():
if name != "pos_embedding":
param.data = param.data.to(model_dtype)
if model_dtype in [torch.float8_e4m3fn, torch.float8_e4m3fn] and fp8_fastmode:
if model_dtype in [torch.float8_e4m3fn, torch.float8_e5m2] and fp8_fastmode:
from ..fp8_optimization import convert_fp8_linear
convert_fp8_linear(self.dit, torch.bfloat16)