apply review suggestions

Author: a-r-r-o-w

src/diffusers/__init__.py

@@ -187,7 +187,6 @@
             "EDMEulerScheduler",
             "EulerAncestralDiscreteScheduler",
             "EulerDiscreteScheduler",
-            "FlowDPMSolverMultistepScheduler",
             "FlowMatchEulerDiscreteScheduler",
             "FlowMatchHeunDiscreteScheduler",
             "HeunDiscreteScheduler",
@@ -692,7 +691,6 @@
             EDMEulerScheduler,
             EulerAncestralDiscreteScheduler,
             EulerDiscreteScheduler,
-            FlowDPMSolverMultistepScheduler,
             FlowMatchEulerDiscreteScheduler,
             FlowMatchHeunDiscreteScheduler,
             HeunDiscreteScheduler,

src/diffusers/models/attention_processor.py

@@ -5446,11 +5446,6 @@ class SanaLinearAttnProcessor2_0:
     Processor for implementing scaled dot-product linear attention.
     """
 
-    def __init__(self, pad_val=1.0, eps=1e-15):
-        self.pad_val = pad_val
-        self.eps = eps
-        self.kernel_func = nn.ReLU(inplace=False)
-
     def __call__(
         self,
         attn: Attention,
@@ -5471,16 +5466,16 @@ def __call__(
         key = key.transpose(1, 2).unflatten(1, (attn.heads, -1)).transpose(2, 3)
         value = value.transpose(1, 2).unflatten(1, (attn.heads, -1))
 
-        query = self.kernel_func(query)
-        key = self.kernel_func(key)
+        query = F.relu(query)
+        key = F.relu(key)
 
         query, key, value = query.float(), key.float(), value.float()
 
-        value = F.pad(value, (0, 0, 0, 1), mode="constant", value=self.pad_val)
+        value = F.pad(value, (0, 0, 0, 1), mode="constant", value=1.0)
         scores = torch.matmul(value, key)
         hidden_states = torch.matmul(scores, query)
 
-        hidden_states = hidden_states[:, :, :-1] / (hidden_states[:, :, -1:] + self.eps)
+        hidden_states = hidden_states[:, :, :-1] / (hidden_states[:, :, -1:] + 1e-15)
         hidden_states = hidden_states.flatten(1, 2).transpose(1, 2)
         hidden_states = hidden_states.to(original_dtype)
 
@@ -5498,11 +5493,6 @@ class PAGCFGSanaLinearAttnProcessor2_0:
     Processor for implementing scaled dot-product linear attention.
     """
 
-    def __init__(self, pad_val=1.0, eps=1e-15):
-        self.pad_val = pad_val
-        self.eps = eps
-        self.kernel_func = nn.ReLU(inplace=False)
-
     def __call__(
         self,
         attn: Attention,
@@ -5523,16 +5513,16 @@ def __call__(
         key = key.transpose(1, 2).unflatten(1, (attn.heads, -1)).transpose(2, 3)
         value = value.transpose(1, 2).unflatten(1, (attn.heads, -1))
 
-        query = self.kernel_func(query)
-        key = self.kernel_func(key)
+        query = F.relu(query)
+        key = F.relu(key)
 
         query, key, value = query.float(), key.float(), value.float()
 
-        value = F.pad(value, (0, 0, 0, 1), mode="constant", value=self.pad_val)
+        value = F.pad(value, (0, 0, 0, 1), mode="constant", value=1.0)
         scores = torch.matmul(value, key)
         hidden_states_org = torch.matmul(scores, query)
 
-        hidden_states_org = hidden_states_org[:, :, :-1] / (hidden_states_org[:, :, -1:] + self.eps)
+        hidden_states_org = hidden_states_org[:, :, :-1] / (hidden_states_org[:, :, -1:] + 1e-15)
         hidden_states_org = hidden_states_org.flatten(1, 2).transpose(1, 2)
         hidden_states_org = hidden_states_org.to(original_dtype)
 
@@ -5558,11 +5548,6 @@ class PAGIdentitySanaLinearAttnProcessor2_0:
     Processor for implementing scaled dot-product linear attention.
     """
 
-    def __init__(self, pad_val=1.0, eps=1e-15):
-        self.pad_val = pad_val
-        self.eps = eps
-        self.kernel_func = nn.ReLU(inplace=False)
-
     def __call__(
         self,
         attn: Attention,
@@ -5587,14 +5572,14 @@ def __call__(
 
         query, key, value = query.float(), key.float(), value.float()
 
-        value = F.pad(value, (0, 0, 0, 1), mode="constant", value=self.pad_val)
+        value = F.pad(value, (0, 0, 0, 1), mode="constant", value=1.0)
         scores = torch.matmul(value, key)
         hidden_states_org = torch.matmul(scores, query)
 
         if hidden_states_org.dtype in [torch.float16, torch.bfloat16]:
             hidden_states_org = hidden_states_org.float()
 
-        hidden_states_org = hidden_states_org[:, :, :-1] / (hidden_states_org[:, :, -1:] + self.eps)
+        hidden_states_org = hidden_states_org[:, :, :-1] / (hidden_states_org[:, :, -1:] + 1e-15)
         hidden_states_org = hidden_states_org.flatten(1, 2).transpose(1, 2)
         hidden_states_org = hidden_states_org.to(original_dtype)
 

src/diffusers/pipelines/pag/pipeline_pag_sana.py

@@ -25,7 +25,7 @@
 from ...image_processor import PixArtImageProcessor
 from ...models import AutoencoderDC, SanaTransformer2DModel
 from ...models.attention_processor import PAGCFGSanaLinearAttnProcessor2_0, PAGIdentitySanaLinearAttnProcessor2_0
-from ...schedulers import FlowDPMSolverMultistepScheduler
+from ...schedulers import FlowMatchEulerDiscreteScheduler
 from ...utils import (
     BACKENDS_MAPPING,
     is_bs4_available,
@@ -140,7 +140,7 @@ def __init__(
         text_encoder: AutoModelForCausalLM,
         vae: AutoencoderDC,
         transformer: SanaTransformer2DModel,
-        scheduler: FlowDPMSolverMultistepScheduler,
+        scheduler: FlowMatchEulerDiscreteScheduler,
         pag_applied_layers: Union[str, List[str]] = "transformer_blocks.0",
     ):
         super().__init__()
@@ -316,7 +316,7 @@ def check_inputs(
         prompt_attention_mask=None,
         negative_prompt_attention_mask=None,
     ):
-        if height % 8 != 0 or width % 8 != 0:
+        if height % 32 != 0 or width % 32 != 0:
             raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
 
         if callback_on_step_end_tensor_inputs is not None and not all(

src/diffusers/pipelines/sana/pipeline_sana.py

@@ -24,7 +24,7 @@
 from ...callbacks import MultiPipelineCallbacks, PipelineCallback
 from ...image_processor import PixArtImageProcessor
 from ...models import AutoencoderDC, SanaTransformer2DModel
-from ...schedulers import FlowDPMSolverMultistepScheduler
+from ...schedulers import FlowMatchEulerDiscreteScheduler
 from ...utils import (
     BACKENDS_MAPPING,
     is_bs4_available,
@@ -137,7 +137,7 @@ def __init__(
         text_encoder: AutoModelForCausalLM,
         vae: AutoencoderDC,
         transformer: SanaTransformer2DModel,
-        scheduler: FlowDPMSolverMultistepScheduler,
+        scheduler: FlowMatchEulerDiscreteScheduler,
     ):
         super().__init__()
 
@@ -187,8 +187,7 @@ def encode_prompt(
                 Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
                 provided, text embeddings will be generated from `prompt` input argument.
             negative_prompt_embeds (`torch.Tensor`, *optional*):
-                Pre-generated negative text embeddings. For Sana, it's should be the embeddings of the ""
-                string.
+                Pre-generated negative text embeddings. For Sana, it's should be the embeddings of the "" string.
             clean_caption (`bool`, defaults to `False`):
                 If `True`, the function will preprocess and clean the provided caption before encoding.
             max_sequence_length (`int`, defaults to 300): Maximum sequence length to use for the prompt.
@@ -325,7 +324,7 @@ def check_inputs(
         prompt_attention_mask=None,
         negative_prompt_attention_mask=None,
     ):
-        if height % 8 != 0 or width % 8 != 0:
+        if height % 32 != 0 or width % 32 != 0:
             raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
 
         if callback_on_step_end_tensor_inputs is not None and not all(