Compare commits

...

2 Commits

Author SHA1 Message Date
Sayak Paul
e1ea8554da Merge branch 'main' into add-attentionmixin-qwen-image 2025-08-22 21:04:20 +05:30
sayakpaul
d15ecd2995 add attentionmixin to qwen image 2025-08-22 20:57:00 +05:30

View File

@@ -24,7 +24,7 @@ from ...configuration_utils import ConfigMixin, register_to_config
from ...loaders import FromOriginalModelMixin, PeftAdapterMixin
from ...utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers
from ...utils.torch_utils import maybe_allow_in_graph
from ..attention import FeedForward
from ..attention import AttentionMixin, FeedForward
from ..attention_dispatch import dispatch_attention_fn
from ..attention_processor import Attention
from ..cache_utils import CacheMixin
@@ -469,7 +469,9 @@ class QwenImageTransformerBlock(nn.Module):
return encoder_hidden_states, hidden_states
class QwenImageTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin, CacheMixin):
class QwenImageTransformer2DModel(
ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin, CacheMixin, AttentionMixin
):
"""
The Transformer model introduced in Qwen.