|
4 | 4 |
|
5 | 5 | from comfy.ldm.modules.attention import optimized_attention |
6 | 6 | import comfy.model_management |
7 | | -import logging |
| 7 | +import comfy.quant_ops |
8 | 8 |
|
9 | 9 |
|
10 | 10 | def attention(q: Tensor, k: Tensor, v: Tensor, pe: Tensor, mask=None, transformer_options={}) -> Tensor: |
@@ -44,21 +44,15 @@ def _apply_rope(xq: Tensor, xk: Tensor, freqs_cis: Tensor): |
44 | 44 | return apply_rope1(xq, freqs_cis), apply_rope1(xk, freqs_cis) |
45 | 45 |
|
46 | 46 |
|
47 | | -try: |
48 | | - import comfy.quant_ops |
49 | | - q_apply_rope = comfy.quant_ops.ck.apply_rope |
50 | | - q_apply_rope1 = comfy.quant_ops.ck.apply_rope1 |
51 | | - def apply_rope(xq, xk, freqs_cis): |
52 | | - if comfy.model_management.in_training: |
53 | | - return _apply_rope(xq, xk, freqs_cis) |
54 | | - else: |
55 | | - return apply_rope1(xq, freqs_cis), apply_rope1(xk, freqs_cis) |
56 | | - def apply_rope1(x, freqs_cis): |
57 | | - if comfy.model_management.in_training: |
58 | | - return _apply_rope1(x, freqs_cis) |
59 | | - else: |
60 | | - return q_apply_rope1(x, freqs_cis) |
61 | | -except: |
62 | | - logging.warning("No comfy kitchen, using old apply_rope functions.") |
63 | | - apply_rope = _apply_rope |
64 | | - apply_rope1 = _apply_rope1 |
| 47 | +def apply_rope(xq, xk, freqs_cis): |
| 48 | + if comfy.model_management.in_training: |
| 49 | + return _apply_rope(xq, xk, freqs_cis) |
| 50 | + else: |
| 51 | + return comfy.quant_ops.ck.apply_rope(xq, xk, freqs_cis) |
| 52 | + |
| 53 | + |
| 54 | +def apply_rope1(x, freqs_cis): |
| 55 | + if comfy.model_management.in_training: |
| 56 | + return _apply_rope1(x, freqs_cis) |
| 57 | + else: |
| 58 | + return comfy.quant_ops.ck.apply_rope1(x, freqs_cis) |
0 commit comments