github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Experimental features and bug fixes

Browse Source
This commit is contained in:
Jaret Burkett
2025-02-04 13:36:34 -07:00
parent e6180d1e1d
commit 216ab164ce
6 changed files with 26 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
toolkit/models/diffusion_feature_extraction.py
View File

@@ -250,13 +250,14 @@ class DiffusionFeatureExtractor3(nn.Module):
bs = noise_pred.shape[0]
noise_pred_chunks = torch.chunk(noise_pred, bs)
timestep_chunks = torch.chunk(timesteps, bs)
noisy_latent_chunks = torch.chunk(noisy_latents, bs)
stepped_chunks = []
for idx in range(bs):
model_output = noise_pred_chunks[idx]
timestep = timestep_chunks[idx]
scheduler._step_index = None
scheduler._init_step_index(timestep)
sample = noisy_latents.to(torch.float32)
sample = noisy_latent_chunks[idx].to(torch.float32)
sigma = scheduler.sigmas[scheduler.step_index]
sigma_next = scheduler.sigmas[-1] # use last sigma for final step