Bug fixes and little improvements here and there.

2024-06-08 06:24:20 -06:00
parent 833c833f28
commit 3f3636b788
12 changed files with 358 additions and 117 deletions
--- a/toolkit/guidance.py
+++ b/toolkit/guidance.py
@@ -240,7 +240,7 @@ def get_direct_guidance_loss(

    noise_pred_uncond, noise_pred_cond = torch.chunk(prediction, 2, dim=0)

-    guidance_scale = 1.0
+    guidance_scale = 1.25
    guidance_pred = noise_pred_uncond + guidance_scale * (
            noise_pred_cond - noise_pred_uncond
    )
@@ -586,6 +586,8 @@ def get_guided_tnt(

    loss = prior_loss + this_loss - that_loss

+    loss = loss.mean()
+
    loss.backward()

    # detach it so parent class can run backward on no grads without throwing error