update guider: remove distilled guidannce scale, simplify prepare_inputs

Author: yiyixuxu

src/diffusers/guiders/adaptive_projected_guidance.py

@@ -77,19 +77,14 @@ def __init__(
         self.use_original_formulation = use_original_formulation
         self.momentum_buffer = None
 
-    def prepare_inputs(
-        self, data: "BlockState", input_fields: Optional[Dict[str, Union[str, Tuple[str, str]]]] = None
-    ) -> List["BlockState"]:
-        if input_fields is None:
-            input_fields = self._input_fields
-
+    def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) -> List["BlockState"]:
         if self._step == 0:
             if self.adaptive_projected_guidance_momentum is not None:
                 self.momentum_buffer = MomentumBuffer(self.adaptive_projected_guidance_momentum)
         tuple_indices = [0] if self.num_conditions == 1 or not self._is_apg_enabled() else [0, 1]
         data_batches = []
         for tuple_idx, input_prediction in zip(tuple_indices, self._input_predictions):
-            data_batch = self._prepare_batch(input_fields, data, tuple_idx, input_prediction)
+            data_batch = self._prepare_batch(data, tuple_idx, input_prediction)
             data_batches.append(data_batch)
         return data_batches
 

src/diffusers/guiders/adaptive_projected_guidance_mix.py

@@ -75,7 +75,6 @@ def __init__(
         stop: float = 1.0,
         adaptive_projected_guidance_start_step: int = 5,
         enabled: bool = True,
-        distilled_guidance_scale: Optional[float] = None,
     ):
         super().__init__(start, stop, enabled)
 
@@ -88,13 +87,9 @@ def __init__(
         self.adaptive_projected_guidance_start_step = adaptive_projected_guidance_start_step
         self.use_original_formulation = use_original_formulation
         self.momentum_buffer = None
-        self.distilled_guidance_scale = distilled_guidance_scale
 
     def prepare_inputs(
-        self, data: "BlockState", input_fields: Optional[Dict[str, Union[str, Tuple[str, str]]]] = None
-    ) -> List["BlockState"]:
-        if input_fields is None:
-            input_fields = self._input_fields
+        self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) -> List["BlockState"]:
 
         if self._step == 0:
             if self.adaptive_projected_guidance_momentum is not None:
@@ -104,7 +99,7 @@ def prepare_inputs(
         )
         data_batches = []
         for tuple_idx, input_prediction in zip(tuple_indices, self._input_predictions):
-            data_batch = self._prepare_batch(input_fields, data, tuple_idx, input_prediction)
+            data_batch = self._prepare_batch(data, tuple_idx, input_prediction)
             data_batches.append(data_batch)
         return data_batches
 

src/diffusers/guiders/auto_guidance.py

@@ -133,16 +133,11 @@ def cleanup_models(self, denoiser: torch.nn.Module) -> None:
                 registry = HookRegistry.check_if_exists_or_initialize(denoiser)
                 registry.remove_hook(name, recurse=True)
 
-    def prepare_inputs(
-        self, data: "BlockState", input_fields: Optional[Dict[str, Union[str, Tuple[str, str]]]] = None
-    ) -> List["BlockState"]:
-        if input_fields is None:
-            input_fields = self._input_fields
-
+    def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) -> List["BlockState"]:
         tuple_indices = [0] if self.num_conditions == 1 or not self._is_ag_enabled() else [0, 1]
         data_batches = []
         for tuple_idx, input_prediction in zip(tuple_indices, self._input_predictions):
-            data_batch = self._prepare_batch(input_fields, data, tuple_idx, input_prediction)
+            data_batch = self._prepare_batch(data, tuple_idx, input_prediction)
             data_batches.append(data_batch)
         return data_batches
 

src/diffusers/guiders/classifier_free_guidance.py

@@ -52,20 +52,10 @@ class ClassifierFreeGuidance(BaseGuidance):
 
     Use `use_original_formulation=True` to switch to the original formulation.
 
-    **Guidance-Distilled Models:**
-
-    For models with distilled guidance (guidance baked into the model via distillation), set `distilled_guidance_scale`
-    to the desired guidance value. The pipeline will pass this to the model during forward passes. Set to `None` for
-    regular (non-distilled) models.
-
     Args:
         guidance_scale (`float`, defaults to `7.5`):
             CFG scale applied by this guider during post-processing. Higher values = stronger prompt conditioning but
             may reduce quality. Typical range: 1.0-20.0.
-        distilled_guidance_scale (`float`, *optional*, defaults to `None`):
-            Guidance scale for distilled models, passed directly to the model during forward pass. If `None`, assumes a
-            regular (non-distilled) model. Allows pipelines to configure different defaults for distilled vs.
-            non-distilled models. Typical range for distilled models: 1.0-8.0.
         guidance_rescale (`float`, defaults to `0.0`):
             Rescaling factor to prevent overexposure from high guidance scales. Based on [Common Diffusion Noise
             Schedules and Sample Steps are Flawed](https://huggingface.co/papers/2305.08891). Range: 0.0 (no rescaling)
@@ -89,7 +79,6 @@ class ClassifierFreeGuidance(BaseGuidance):
     def __init__(
         self,
         guidance_scale: float = 7.5,
-        distilled_guidance_scale: Optional[float] = None,
         guidance_rescale: float = 0.0,
         use_original_formulation: bool = False,
         start: float = 0.0,
@@ -99,20 +88,15 @@ def __init__(
         super().__init__(start, stop, enabled)
 
         self.guidance_scale = guidance_scale
-        self.distilled_guidance_scale = distilled_guidance_scale
         self.guidance_rescale = guidance_rescale
         self.use_original_formulation = use_original_formulation
 
-    def prepare_inputs(
-        self, data: "BlockState", input_fields: Optional[Dict[str, Union[str, Tuple[str, str]]]] = None
-    ) -> List["BlockState"]:
-        if input_fields is None:
-            input_fields = self._input_fields
+    def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) -> List["BlockState"]:
 
         tuple_indices = [0] if self.num_conditions == 1 or not self._is_cfg_enabled() else [0, 1]
         data_batches = []
         for tuple_idx, input_prediction in zip(tuple_indices, self._input_predictions):
-            data_batch = self._prepare_batch(input_fields, data, tuple_idx, input_prediction)
+            data_batch = self._prepare_batch(data, tuple_idx, input_prediction)
             data_batches.append(data_batch)
         return data_batches
 

src/diffusers/guiders/classifier_free_zero_star_guidance.py

@@ -77,16 +77,11 @@ def __init__(
         self.guidance_rescale = guidance_rescale
         self.use_original_formulation = use_original_formulation
 
-    def prepare_inputs(
-        self, data: "BlockState", input_fields: Optional[Dict[str, Union[str, Tuple[str, str]]]] = None
-    ) -> List["BlockState"]:
-        if input_fields is None:
-            input_fields = self._input_fields
-
+    def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) -> List["BlockState"]:
         tuple_indices = [0] if self.num_conditions == 1 or not self._is_cfg_enabled() else [0, 1]
         data_batches = []
         for tuple_idx, input_prediction in zip(tuple_indices, self._input_predictions):
-            data_batch = self._prepare_batch(input_fields, data, tuple_idx, input_prediction)
+            data_batch = self._prepare_batch(data, tuple_idx, input_prediction)
             data_batches.append(data_batch)
         return data_batches
 

src/diffusers/guiders/frequency_decoupled_guidance.py

@@ -218,16 +218,11 @@ def __init__(
                 f"({len(self.guidance_scales)})"
             )
 
-    def prepare_inputs(
-        self, data: "BlockState", input_fields: Optional[Dict[str, Union[str, Tuple[str, str]]]] = None
-    ) -> List["BlockState"]:
-        if input_fields is None:
-            input_fields = self._input_fields
-
+    def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) -> List["BlockState"]:
         tuple_indices = [0] if self.num_conditions == 1 or not self._is_fdg_enabled() else [0, 1]
         data_batches = []
         for tuple_idx, input_prediction in zip(tuple_indices, self._input_predictions):
-            data_batch = self._prepare_batch(input_fields, data, tuple_idx, input_prediction)
+            data_batch = self._prepare_batch(data, tuple_idx, input_prediction)
             data_batches.append(data_batch)
         return data_batches
 

src/diffusers/guiders/guider_utils.py

@@ -190,8 +190,7 @@ def num_conditions(self) -> int:
     @classmethod
     def _prepare_batch(
         cls,
-        input_fields: Dict[str, Union[str, Tuple[str, str]]],
-        data: "BlockState",
+        data: Dict[str, Tuple[torch.Tensor, torch.Tensor]],
         tuple_index: int,
         identifier: str,
     ) -> "BlockState":
@@ -217,24 +216,17 @@ def _prepare_batch(
         """
         from ..modular_pipelines.modular_pipeline import BlockState
 
-        if isinstance(data, dict):
-            data = BlockState(**data)
 
-        if input_fields is None:
-            raise ValueError(
-                "Input fields cannot be None. Please pass `input_fields` to `prepare_inputs` or call `set_input_fields` before preparing inputs."
-            )
         data_batch = {}
-        for key, value in input_fields.items():
+        for key, value in data.items():
             try:
-                if isinstance(value, str):
-                    data_batch[key] = getattr(data, value)
+                if isinstance(value, torch.Tensor):
+                    data_batch[key] = value
                 elif isinstance(value, tuple):
-                    data_batch[key] = getattr(data, value[tuple_index])
+                    data_batch[key] = value[tuple_index]
                 else:
-                    # We've already checked that value is a string or a tuple of strings with length 2
-                    pass
-            except AttributeError:
+                    raise ValueError(f"Invalid value type: {type(value)}")
+            except ValueError:
                 logger.debug(f"`data` does not have attribute(s) {value}, skipping.")
         data_batch[cls._identifier_key] = identifier
         return BlockState(**data_batch)

src/diffusers/guiders/perturbed_attention_guidance.py

@@ -169,12 +169,7 @@ def cleanup_models(self, denoiser: torch.nn.Module) -> None:
                 registry.remove_hook(hook_name, recurse=True)
 
     # Copied from diffusers.guiders.skip_layer_guidance.SkipLayerGuidance.prepare_inputs
-    def prepare_inputs(
-        self, data: "BlockState", input_fields: Optional[Dict[str, Union[str, Tuple[str, str]]]] = None
-    ) -> List["BlockState"]:
-        if input_fields is None:
-            input_fields = self._input_fields
-
+    def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) -> List["BlockState"]:
         if self.num_conditions == 1 or not self._is_cfg_enabled() and not self._is_slg_enabled():
             tuple_indices = [0]
             input_predictions = ["pred_cond"]
@@ -188,7 +183,7 @@ def prepare_inputs(
             input_predictions = ["pred_cond", "pred_uncond", "pred_cond_skip"]
         data_batches = []
         for tuple_idx, input_prediction in zip(tuple_indices, input_predictions):
-            data_batch = self._prepare_batch(input_fields, data, tuple_idx, input_prediction)
+            data_batch = self._prepare_batch(data, tuple_idx, input_prediction)
             data_batches.append(data_batch)
         return data_batches
 

src/diffusers/guiders/skip_layer_guidance.py

@@ -165,12 +165,7 @@ def cleanup_models(self, denoiser: torch.nn.Module) -> None:
             for hook_name in self._skip_layer_hook_names:
                 registry.remove_hook(hook_name, recurse=True)
 
-    def prepare_inputs(
-        self, data: "BlockState", input_fields: Optional[Dict[str, Union[str, Tuple[str, str]]]] = None
-    ) -> List["BlockState"]:
-        if input_fields is None:
-            input_fields = self._input_fields
-
+    def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) -> List["BlockState"]:
         if self.num_conditions == 1 or not self._is_cfg_enabled() and not self._is_slg_enabled():
             tuple_indices = [0]
             input_predictions = ["pred_cond"]
@@ -184,7 +179,7 @@ def prepare_inputs(
             input_predictions = ["pred_cond", "pred_uncond", "pred_cond_skip"]
         data_batches = []
         for tuple_idx, input_prediction in zip(tuple_indices, input_predictions):
-            data_batch = self._prepare_batch(input_fields, data, tuple_idx, input_prediction)
+            data_batch = self._prepare_batch(data, tuple_idx, input_prediction)
             data_batches.append(data_batch)
         return data_batches
 

src/diffusers/guiders/smoothed_energy_guidance.py

@@ -154,12 +154,7 @@ def cleanup_models(self, denoiser: torch.nn.Module):
             for hook_name in self._seg_layer_hook_names:
                 registry.remove_hook(hook_name, recurse=True)
 
-    def prepare_inputs(
-        self, data: "BlockState", input_fields: Optional[Dict[str, Union[str, Tuple[str, str]]]] = None
-    ) -> List["BlockState"]:
-        if input_fields is None:
-            input_fields = self._input_fields
-
+    def prepare_inputs(self, data: Dict[str, Tuple[torch.Tensor, torch.Tensor]]) -> List["BlockState"]:
         if self.num_conditions == 1 or not self._is_cfg_enabled() and not self._is_seg_enabled():
             tuple_indices = [0]
             input_predictions = ["pred_cond"]
@@ -173,7 +168,7 @@ def prepare_inputs(
             input_predictions = ["pred_cond", "pred_uncond", "pred_cond_seg"]
         data_batches = []
         for tuple_idx, input_prediction in zip(tuple_indices, input_predictions):
-            data_batch = self._prepare_batch(input_fields, data, tuple_idx, input_prediction)
+            data_batch = self._prepare_batch(data, tuple_idx, input_prediction)
             data_batches.append(data_batch)
         return data_batches