ostris/ai-toolkit
1
0
Fork 0
mirror of https://github.com/ostris/ai-toolkit.git synced 2026-05-01 03:31:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

Added flux training instructions

Browse Source
This commit is contained in:
Jaret Burkett
2024-08-10 14:10:02 -06:00
parent b3e03295ad
commit 2308ef2868
4 changed files with 128 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
toolkit/pipelines.py
View File

@@ -1349,7 +1349,6 @@ class FluxWithCFGPipeline(FluxPipeline):
noise_pred_text = self.transformer(
hidden_states=latents,
# YiYi notes: divide it by 1000 for now because we scale it by 1000 in the transforme rmodel (we should not keep it but I want to keep the inputs same for the model for testing)
timestep=timestep / 1000,
guidance=guidance,
pooled_projections=pooled_prompt_embeds,
@@ -1363,7 +1362,6 @@ class FluxWithCFGPipeline(FluxPipeline):
# todo combine these
noise_pred_uncond = self.transformer(
hidden_states=latents,
# YiYi notes: divide it by 1000 for now because we scale it by 1000 in the transforme rmodel (we should not keep it but I want to keep the inputs same for the model for testing)
timestep=timestep / 1000,
guidance=guidance,
pooled_projections=negative_pooled_prompt_embeds,
@@ -1376,8 +1374,6 @@ class FluxWithCFGPipeline(FluxPipeline):
noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)
# compute the previous noisy sample x_t -> x_t-1
latents_dtype = latents.dtype
latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]