Skip to content

Commit

Permalink
Fast previews for mochi.
Browse files Browse the repository at this point in the history
  • Loading branch information
comfyanonymous committed Nov 10, 2024
1 parent 9c1ed58 commit bdeb1c1
Show file tree
Hide file tree
Showing 2 changed files with 21 additions and 2 deletions.
16 changes: 15 additions & 1 deletion comfy/latent_formats.py
Original file line number Diff line number Diff line change
Expand Up @@ -190,7 +190,21 @@ def __init__(self):
0.9294154431013696, 1.3720942357788521, 0.881393668867029,
0.9168315692124348, 0.9185249279345552, 0.9274757570805041]).view(1, self.latent_channels, 1, 1, 1)

self.latent_rgb_factors = None #TODO
self.latent_rgb_factors =[
[-0.0069, -0.0045, 0.0018],
[ 0.0154, -0.0692, -0.0274],
[ 0.0333, 0.0019, 0.0206],
[-0.1390, 0.0628, 0.1678],
[-0.0725, 0.0134, -0.1898],
[ 0.0074, -0.0270, -0.0209],
[-0.0176, -0.0277, -0.0221],
[ 0.5294, 0.5204, 0.3852],
[-0.0326, -0.0446, -0.0143],
[-0.0659, 0.0153, -0.0153],
[ 0.0185, -0.0217, 0.0014],
[-0.0396, -0.0495, -0.0281]
]
self.latent_rgb_factors_bias = [-0.0940, -0.1418, -0.1453]
self.taesd_decoder_name = None #TODO

def process_in(self, latent):
Expand Down
7 changes: 6 additions & 1 deletion latent_preview.py
Original file line number Diff line number Diff line change
Expand Up @@ -47,7 +47,12 @@ def decode_latent_to_preview(self, x0):
if self.latent_rgb_factors_bias is not None:
self.latent_rgb_factors_bias = self.latent_rgb_factors_bias.to(dtype=x0.dtype, device=x0.device)

latent_image = torch.nn.functional.linear(x0[0].permute(1, 2, 0), self.latent_rgb_factors, bias=self.latent_rgb_factors_bias)
if x0.ndim == 5:
x0 = x0[0, :, 0]
else:
x0 = x0[0]

latent_image = torch.nn.functional.linear(x0.movedim(0, -1), self.latent_rgb_factors, bias=self.latent_rgb_factors_bias)
# latent_image = x0[0].permute(1, 2, 0) @ self.latent_rgb_factors

return preview_to_image(latent_image)
Expand Down

0 comments on commit bdeb1c1

Please sign in to comment.