Fall back to using c_crossattn if y key does not exist

2024-03-24 04:01:18 -06:00
parent 9a6d6445bc
commit 9b2a19eef9
1 changed files with 3 additions and 2 deletions
@@ -94,11 +94,12 @@ class CADS:
            c = args["c"]

            if noise_scale > 0.0:
+                apply_to = c.get(key, c["c_crossattn"])
                gamma = cads_gamma(timestep)
-                for i in range(c[key].size(dim=0)):
+                for i in range(apply_to.size(dim=0)):
                    if cond_or_uncond[i % len(cond_or_uncond)] == skip:
                        continue
-                    c[key][i] = cads_noise(gamma, c[key][i])
+                    apply_to[i] = cads_noise(gamma, apply_to[i])

            if previous_wrapper:
                return previous_wrapper(apply_model, args)