Set APG double to float for MPS

docs/workflows/inference.md

@@ -218,10 +218,6 @@ model.set_lora_strength(0.0)              # Disable without unloading
 # With multiple LoRAs, target by index:
 model.set_lora_strength(1.0, lora_index=0)
 model.set_lora_strength(0.3, lora_index=1)
-
-# Target only the Diffusion Transformer backbone or conditioner independently:
-model.set_lora_strength(1.0, target="dit")
-model.set_lora_strength(0.0, target="conditioner")
 ```
 
 For full details on LoRA training see [LoRA Training](lora.md).

pyproject.toml

@@ -66,4 +66,5 @@ exclude = [
     "stable_audio_3/interface",
     "stable_audio_3/data",
     "stable_audio_3/training",
+    "optimized",
 ]

stable_audio_3/model.py

@@ -28,7 +28,7 @@ def __init__(self, model, model_config, device, model_half):
         torch.backends.cudnn.benchmark = False
 
     @staticmethod
-    def from_pretrained(model_name_or_path, device=None, model_half=True):
+    def from_pretrained(model_name, device=None, model_half=True):
         # Load the model and any necessary components here
         if device is None and torch.cuda.is_available():
             device = "cuda"
@@ -38,18 +38,18 @@ def from_pretrained(model_name_or_path, device=None, model_half=True):
             device = "cpu"
 
         if not torch.cuda.is_available():
-            if model_name_or_path in ("medium", "medium-base"):
+            if model_name in ("medium", "medium-base"):
                 print(
-                    f"Warning: You are loading the {model_name_or_path} model without a GPU. This model is not designed to run on cpu"
+                    f"Warning: You are loading the {model_name} model without a GPU. This model is not designed to run on cpu"
                 )
             model_half = False
 
-        if model_name_or_path not in all_models:
+        if model_name not in all_models:
             raise ValueError(
-                f"Unknown model '{model_name_or_path}'. Valid models: {list(all_models)}"
+                f"Unknown model '{model_name}'. Valid models: {list(all_models)}"
             )
 
-        model_cfg = all_models[model_name_or_path]
+        model_cfg = all_models[model_name]
         local_config, local_ckpt = model_cfg.resolve()
         with open(local_config) as f:
             model_config = json.load(f)

stable_audio_3/models/dit.py

@@ -319,11 +319,11 @@ def apg_project(self, v0, v1, padding_mask=None):
                           If provided, only valid positions contribute to the projection.
         """
         dtype = v0.dtype
-        v0, v1 = v0.double(), v1.double()
+        v0, v1 = v0.float(), v1.float()
 
         if padding_mask is not None:
             # Expand mask to match tensor shape: (B, T) -> (B, 1, T)
-            mask = padding_mask.unsqueeze(1).double()
+            mask = padding_mask.unsqueeze(1).float()
             # Zero out padding positions for projection computation
             v0_masked = v0 * mask
             v1_masked = v1 * mask