nicolabs/Open-Sora
1
0
Fork 0
mirror of https://github.com/hpcaitech/Open-Sora.git synced 2026-04-10 12:49:38 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

Force fp16 input to fp32 to avoid nan output in timestep_transform

Browse source
This commit is contained in:
HangXu 2024-06-21 11:15:39 +03:00 committed by GitHub
parent 033c2b3c82
commit 04d2ee0182
1 changed files with 5 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
opensora/schedulers/rf/rectified_flow.py
View file

@ -15,6 +15,11 @@ def timestep_transform(
scale=1.0,
num_timesteps=1,
):
# Force fp16 input to fp32 to avoid nan output
for key in ["height", "width", "num_frames"]:
if model_kwargs[key].dtype == torch.float16:
model_kwargs[key] = model_kwargs[key].float()
t = t / num_timesteps
resolution = model_kwargs["height"] * model_kwargs["width"]
ratio_space = (resolution / base_resolution).sqrt()