ostris/ai-toolkit
1
0
Fork 0
mirror of https://github.com/ostris/ai-toolkit.git synced 2026-03-13 22:49:48 +00:00
Code Issues Packages Projects Releases Wiki Activity

Added initial support for f-lite model

Browse Source
This commit is contained in:
Jaret Burkett
2025-05-01 11:15:18 -06:00
parent 5890e67a46
commit d9700bdb99
9 changed files with 1076 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
toolkit/models/diffusion_feature_extraction.py
View File

@@ -249,7 +249,8 @@ class DiffusionFeatureExtractor3(nn.Module):
# lpips_weight=1.0,
lpips_weight=10.0,
clip_weight=0.1,
pixel_weight=0.1
pixel_weight=0.1,
model=None
):
dtype = torch.bfloat16
device = self.vae.device
@@ -274,7 +275,10 @@ class DiffusionFeatureExtractor3(nn.Module):
# stepped_latents = torch.cat(stepped_chunks, dim=0)
stepped_latents = noise - noise_pred
if model is not None and hasattr(model, 'get_stepped_pred'):
stepped_latents = model.get_stepped_pred(noise_pred, noise)
else:
stepped_latents = noise - noise_pred
latents = stepped_latents.to(self.vae.device, dtype=self.vae.dtype)