RIFLEx support

This commit is contained in:
DeepBeepMeep 2025-03-02 17:55:30 +01:00
parent 7a8dcbf63d
commit f3b365a5da

View File

@ -677,16 +677,14 @@ class WanVAE:
""" """
videos: A list of videos each with shape [C, T, H, W]. videos: A list of videos each with shape [C, T, H, W].
""" """
with amp.autocast(dtype=self.dtype): return [
return [ self.model.encode(u.unsqueeze(0), self.scale).float().squeeze(0)
self.model.encode(u.unsqueeze(0), self.scale).float().squeeze(0) for u in videos
for u in videos ]
]
def decode(self, zs): def decode(self, zs):
with amp.autocast(dtype=self.dtype): return [
return [ self.model.decode(u.unsqueeze(0),
self.model.decode(u.unsqueeze(0), self.scale).float().clamp_(-1, 1).squeeze(0)
self.scale).float().clamp_(-1, 1).squeeze(0) for u in zs
for u in zs ]
]