Skip to content

Commit c3a3234

Browse files
committed
Fix-3
Signed-off-by: Amit Raj <[email protected]>
1 parent 2779920 commit c3a3234

File tree

14 files changed

+400
-296
lines changed

14 files changed

+400
-296
lines changed

QEfficient/diffusers/models/attention.py

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,6 @@
66
# ----------------------------------------------------------------------------
77

88
import torch
9-
109
from diffusers.models.attention import JointTransformerBlock, _chunked_feed_forward
1110

1211

QEfficient/diffusers/models/attention_processor.py

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,6 @@
88
from typing import Optional
99

1010
import torch
11-
1211
from diffusers.models.attention_processor import Attention, JointAttnProcessor2_0
1312

1413

@@ -123,7 +122,6 @@ def __call__(
123122
hidden_states = torch.bmm(attention_probs, value)
124123
else: # self-attention, use blocked attention
125124
# QKV done with block-attention (a la FlashAttentionV2)
126-
print(f"{query.shape = }, {key.shape = }, {value.shape = }")
127125
query_block_size = self.query_block_size
128126
query_seq_len = query.size(-2)
129127
num_blocks = (query_seq_len + query_block_size - 1) // query_block_size

QEfficient/diffusers/models/autoencoders/autoencoder_kl.py

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,6 @@
66
# ----------------------------------------------------------------------------
77

88
import torch
9-
109
from diffusers import AutoencoderKL
1110

1211

QEfficient/diffusers/models/pytorch_transforms.py

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,11 +6,11 @@
66
# -----------------------------------------------------------------------------
77
from typing import Tuple
88

9-
from torch import nn
10-
119
from diffusers.models.attention import JointTransformerBlock
1210
from diffusers.models.attention_processor import Attention, JointAttnProcessor2_0
1311
from diffusers.models.normalization import RMSNorm
12+
from torch import nn
13+
1414
from QEfficient.base.pytorch_transforms import ModuleMappingTransform
1515
from QEfficient.customop.rms_norm import CustomRMSNormAIC
1616
from QEfficient.diffusers.models.attention import QEffJointTransformerBlock
@@ -23,6 +23,11 @@
2323
class CustomOpsTransform(ModuleMappingTransform):
2424
_module_mapping = {RMSNorm: CustomRMSNormAIC}
2525

26+
@classmethod
27+
def apply(cls, model: nn.Module) -> Tuple[nn.Module, bool]:
28+
model, transformed = super().apply(model)
29+
return model, transformed
30+
2631

2732
class AttentionTransform(ModuleMappingTransform):
2833
_module_mapping = {

QEfficient/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,9 +10,9 @@
1010

1111
import numpy as np
1212
import torch
13-
1413
from diffusers import StableDiffusionPipeline
1514
from diffusers.image_processor import VaeImageProcessor
15+
1616
from QEfficient.diffusers.pipelines.pipeline_utils import QEffSafetyChecker, QEffTextEncoder, QEffUNet, QEffVAE
1717
from QEfficient.generation.cloud_infer import QAICInferenceSession
1818
from QEfficient.utils import constants

0 commit comments

Comments
 (0)