Modify env variable

Signed-off-by: lijian <lijian6@sugon.com>

src/diffusers/models/attention.py

@@ -170,8 +170,8 @@ class JointTransformerBlock(nn.Module):
             )
 
         # Attention.
-        use_xformers = os.getenv('USE_XFORMERS', '0')
-        if use_xformers == '1':
+        sd3_use_xformers = os.getenv('SD3_USE_XFORMERS', '0')
+        if sd3_use_xformers == '1':
             self.attn.set_processor(self.processor)
         attn_output, context_attn_output = self.attn(
             hidden_states=norm_hidden_states, encoder_hidden_states=norm_encoder_hidden_states

src/diffusers/models/attention_processor.py

@@ -1129,8 +1129,8 @@ class JointAttnProcessor2_0:
         key = torch.cat([key, encoder_hidden_states_key_proj], dim=1)
         value = torch.cat([value, encoder_hidden_states_value_proj], dim=1)
 
-        use_xformers = os.getenv('USE_XFORMERS', '0')
-        if use_xformers == '1':
+        sd3_use_xformers = os.getenv('SD3_USE_XFORMERS', '0')
+        if sd3_use_xformers == '1':
             query = attn.head_to_batch_dim(query).contiguous()
             key = attn.head_to_batch_dim(key).contiguous()
             value = attn.head_to_batch_dim(value).contiguous()
@@ -1352,7 +1352,9 @@ class XFormersAttnProcessor:
         *args,
         **kwargs,
     ) -> torch.Tensor:
-        from .embeddings import apply_rotary_emb
+        hy_use_xformers = os.getenv('HY_USE_XFORMERS', '0')
+        if hy_use_xformers == '1':
+            from .embeddings import apply_rotary_emb
         if len(args) > 0 or kwargs.get("scale", None) is not None:
             deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`."
             deprecate("scale", "1.0.0", deprecation_message)
@@ -1396,24 +1398,25 @@ class XFormersAttnProcessor:
         key = attn.to_k(encoder_hidden_states)
         value = attn.to_v(encoder_hidden_states)
 
-        ### Add for HunYuanDit xformers ###
-        q_seq_len = query.shape[1]
-        k_seq_len = key.shape[1]
-        inner_dim = key.shape[-1]
-        head_dim = inner_dim // attn.heads
-        if image_rotary_emb is not None:
-            query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
-            key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
-            if attn.norm_q is not None:
-                query = attn.norm_q(query)
-            if attn.norm_k is not None:
-                key = attn.norm_k(key)
-            query = apply_rotary_emb(query, image_rotary_emb)
-            if not attn.is_cross_attention:
-                key = apply_rotary_emb(key, image_rotary_emb)
-            query = query.transpose(1, 2).contiguous().view(batch_size, q_seq_len, -1)
-            key = key.transpose(1, 2).contiguous().view(batch_size, k_seq_len, -1)
-        ### End add ###
+        if hy_use_xformers == '1':
+            ### Add for HunYuanDit xformers ###
+            q_seq_len = query.shape[1]
+            k_seq_len = key.shape[1]
+            inner_dim = key.shape[-1]
+            head_dim = inner_dim // attn.heads
+            if image_rotary_emb is not None:
+                query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+                key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+                if attn.norm_q is not None:
+                    query = attn.norm_q(query)
+                if attn.norm_k is not None:
+                    key = attn.norm_k(key)
+                query = apply_rotary_emb(query, image_rotary_emb)
+                if not attn.is_cross_attention:
+                    key = apply_rotary_emb(key, image_rotary_emb)
+                query = query.transpose(1, 2).contiguous().view(batch_size, q_seq_len, -1)
+                key = key.transpose(1, 2).contiguous().view(batch_size, k_seq_len, -1)
+            ### End add ###
 
         query = attn.head_to_batch_dim(query).contiguous()
         key = attn.head_to_batch_dim(key).contiguous()

src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py

@@ -868,8 +868,8 @@ class StableDiffusion3Pipeline(DiffusionPipeline, SD3LoraLoaderMixin, FromSingle
                 if XLA_AVAILABLE:
                     xm.mark_step()
 
-        use_xformers = os.getenv('USE_XFORMERS', '0')
-        if use_xformers == '1':
+        sd3_use_xformers = os.getenv('SD3_USE_XFORMERS', '0')
+        if sd3_use_xformers == '1':
             self.disable_xformers_memory_efficient_attention()
         if output_type == "latent":
             image = latents