merge PR invoke-ai#495 - keep using float16 in ldm.modules.attention

cognidesign · Sep 11, 2022 · 70aa674 · 70aa674
1 parent 8748370
commit 70aa674
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 3 deletions.
diff --git a/ldm/modules/attention.py b/ldm/modules/attention.py
@@ -181,7 +181,7 @@ def forward(self, x, context=None, mask=None):
         q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q_in, k_in, v_in))
         del q_in, k_in, v_in
 
-        r1 = torch.zeros(q.shape[0], q.shape[1], v.shape[2], device=q.device)
+        r1 = torch.zeros(q.shape[0], q.shape[1], v.shape[2], device=q.device, dtype=q.dtype)
 
         if device_type == 'mps':
             mem_free_total = psutil.virtual_memory().available
@@ -213,7 +213,7 @@ def forward(self, x, context=None, mask=None):
             end = i + slice_size
             s1 = einsum('b i d, b j d -> b i j', q[:, i:end], k) * self.scale
 
-            s2 = s1.softmax(dim=-1)
+            s2 = s1.softmax(dim=-1, dtype=r1.dtype)
             del s1
 
             r1[:, i:end] = einsum('b i j, b j d -> b i d', s2, v)

diff --git a/scripts/dream.py b/scripts/dream.py
@@ -185,7 +185,6 @@ def main_loop(t2i, outdir, prompt_as_dir, parser, infile):
             continue
         if opt.seed is not None and opt.seed < 0:   # retrieve previous value!
             try:
-                print(f'last seeds = {last_seeds}, opt.seed={opt.seed}')
                 opt.seed = last_seeds[opt.seed]
                 print(f'reusing previous seed {opt.seed}')
             except IndexError: