fix

Author: vpj

docs/diffusion/stable_diffusion/model/unet_attention.html

@@ -602,7 +602,7 @@ <h3>Cross Attention Layer</h3>
 <span class="lineno">173</span>        <span class="n">k</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">to_k</span><span class="p">(</span><span class="n">cond</span><span class="p">)</span>
 <span class="lineno">174</span>        <span class="n">v</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">to_v</span><span class="p">(</span><span class="n">cond</span><span class="p">)</span>
 <span class="lineno">175</span>
-<span class="lineno">176</span>        <span class="nb">print</span><span class="p">(</span><span class="s1">&#39;use flash&#39;</span><span class="p">,</span> <span class="n">CrossAttention</span><span class="o">.</span><span class="n">use_flash_attention</span><span class="p">)</span>
+<span class="lineno">176</span>        <span class="nb">print</span><span class="p">(</span><span class="s1">&#39;use flash&#39;</span><span class="p">,</span> <span class="n">CrossAttention</span><span class="o">.</span><span class="n">use_flash_attention</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">flash</span><span class="p">)</span>
 <span class="lineno">177</span>
 <span class="lineno">178</span>        <span class="k">if</span> <span class="n">CrossAttention</span><span class="o">.</span><span class="n">use_flash_attention</span> <span class="ow">and</span> <span class="bp">self</span><span class="o">.</span><span class="n">flash</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">and</span> <span class="n">cond</span> <span class="ow">is</span> <span class="kc">None</span> <span class="ow">and</span> <span class="bp">self</span><span class="o">.</span><span class="n">d_head</span> <span class="o">&lt;=</span> <span class="mi">128</span><span class="p">:</span>
 <span class="lineno">179</span>            <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">flash_attention</span><span class="p">(</span><span class="n">q</span><span class="p">,</span> <span class="n">k</span><span class="p">,</span> <span class="n">v</span><span class="p">)</span>

labml_nn/diffusion/stable_diffusion/model/unet_attention.py

@@ -173,7 +173,7 @@ def forward(self, x: torch.Tensor, cond: Optional[torch.Tensor] = None):
         k = self.to_k(cond)
         v = self.to_v(cond)
 
-        print('use flash', CrossAttention.use_flash_attention)
+        print('use flash', CrossAttention.use_flash_attention, self.flash)
 
         if CrossAttention.use_flash_attention and self.flash is not None and cond is None and self.d_head <= 128:
             return self.flash_attention(q, k, v)