From a6eafe573150ca660171e4df7ebd956c7a43f206 Mon Sep 17 00:00:00 2001
From: Li Xiang <54010254+lixiang007666@users.noreply.github.com>
Date: Mon, 4 Nov 2024 19:04:57 +0800
Subject: [PATCH] Fix diffusers attention register (#1131)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

This PR is done:

- [x] Fix: https://github.com/siliconflow/onediff/issues/1105
- [x] 修复 onediff_quant 的 from_pretrained 兼容 diffusers >= 0.29.1:
https://github.com/siliconflow/onediff-quant/pull/39

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

- **New Features**
- Enhanced attention mechanism to automatically revert to a standard
processor when using specific processor types from the `diffusers`
library.
- Introduced new command-line arguments for the text-to-image script:
`--compile_text_encoder` and `--graph`, allowing for more flexible model
compilation options.

- **Bug Fixes**
- Improved handling of processor types to ensure better compatibility
and functionality within the attention processing framework.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
---
 .../examples/text_to_image_sdxl_enterprise.py          |  1 -
 .../register_diffusers/attention_processor_oflow.py    | 10 ++++++++++
 2 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/onediff_diffusers_extensions/examples/text_to_image_sdxl_enterprise.py b/onediff_diffusers_extensions/examples/text_to_image_sdxl_enterprise.py
index 2c53aa3eb..258c82fa3 100644
--- a/onediff_diffusers_extensions/examples/text_to_image_sdxl_enterprise.py
+++ b/onediff_diffusers_extensions/examples/text_to_image_sdxl_enterprise.py
@@ -83,7 +83,6 @@ def parse_args():
     args.model,
     torch_dtype=torch.float16,
     use_safetensors=True,
-    variant="fp16",
 )
 pipe.to("cuda")
 
diff --git a/src/infer_compiler_registry/register_diffusers/attention_processor_oflow.py b/src/infer_compiler_registry/register_diffusers/attention_processor_oflow.py
index 70684b16c..2aab4f440 100644
--- a/src/infer_compiler_registry/register_diffusers/attention_processor_oflow.py
+++ b/src/infer_compiler_registry/register_diffusers/attention_processor_oflow.py
@@ -361,6 +361,16 @@ def forward(
         # here we simply pass along all tensors to the selected processor class
         # For standard processors that are defined here, `**cross_attention_kwargs` is empty
 
+        from diffusers.models.attention_processor import (
+            AttnProcessor as DiffusersAttnProcessor,
+            AttnProcessor2_0 as DiffusersAttnProcessor2_0,
+        )
+
+        if isinstance(self.processor, DiffusersAttnProcessor) or isinstance(
+            self.processor, DiffusersAttnProcessor2_0
+        ):
+            self.set_processor(AttnProcessor())
+
         return self.processor(
             self,
             hidden_states,