hongshaorou/dreamer4
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

cleanup

Browse Source
This commit is contained in:
lucidrains 2025-10-04 08:04:42 -07:00
parent 0f4783f23c
commit 7cac3d28c5
1 changed files with 4 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
dreamer4/dreamer4.py
View File

@ -1,6 +1,7 @@
from __future__ import annotations from __future__ import annotations
import math import math
from collections import namedtuple
from functools import partial from functools import partial
import torch import torch
@ -49,6 +50,8 @@ except ImportError:
LinearNoBias = partial(Linear, bias = False) LinearNoBias = partial(Linear, bias = False)
TokenizerLosses = namedtuple('TokenizerLosses', ('recon', 'lpips'))
# helpers # helpers
def exists(v): def exists(v):
@ -780,7 +783,7 @@ class VideoTokenizer(Module):
losses = (recon_loss, lpips_loss) losses = (recon_loss, lpips_loss)
return total_loss, losses return total_loss, TokenizerLosses(losses)
# dynamics model, axial space-time transformer # dynamics model, axial space-time transformer