[fix] extend

apg_guidance.py

@@ -17,14 +17,15 @@ def project(
     dims=[-1, -2],
 ):
     dtype = v0.dtype
-    if v0.device.type == "mps":
-        v0, v1 = v0.float(), v1.float()
-    else:
-        v0, v1 = v0.double(), v1.double()
+    device_type = v0.device.type
+    if device_type == "mps":
+        v0, v1 = v0.cpu(), v1.cpu()
+
+    v0, v1 = v0.double(), v1.double()
     v1 = torch.nn.functional.normalize(v1, dim=dims)
     v0_parallel = (v0 * v1).sum(dim=dims, keepdim=True) * v1
     v0_orthogonal = v0 - v0_parallel
-    return v0_parallel.to(dtype), v0_orthogonal.to(dtype)
+    return v0_parallel.to(dtype).to(device_type), v0_orthogonal.to(dtype).to(device_type)
 
 
 def apg_forward(

app.py

@@ -7,7 +7,7 @@ import os
 
 parser = argparse.ArgumentParser()
 parser.add_argument("--checkpoint_path", type=str, default=None)
-parser.add_argument("--server_name", type=str, default="0.0.0.0")
+parser.add_argument("--server_name", type=str, default="127.0.0.1")
 parser.add_argument("--port", type=int, default=7860)
 parser.add_argument("--device_id", type=int, default=0)
 parser.add_argument("--share", action='store_true', default=False)

pipeline_ace_step.py

@@ -68,8 +68,8 @@ class ACEStepPipeline:
         if device.type == "cpu" and torch.backends.mps.is_available():
             device = torch.device("mps")
         self.dtype = torch.bfloat16 if dtype == "bfloat16" else torch.float32
-        if device.type == "mps" and self.dtype == torch.bfloat16:
-            self.dtype = torch.float16
+        if device.type == "mps":
+            self.dtype = torch.float32
         self.device = device
         self.loaded = False
         self.torch_compile = torch_compile
@@ -181,33 +181,33 @@ class ACEStepPipeline:
             last_hidden_states = outputs.last_hidden_state
         attention_mask = inputs["attention_mask"]
         return last_hidden_states, attention_mask
-
+    
     def get_text_embeddings_null(self, texts, device, text_max_length=256, tau=0.01, l_min=8, l_max=10):
         inputs = self.text_tokenizer(texts, return_tensors="pt", padding=True, truncation=True, max_length=text_max_length)
         inputs = {key: value.to(device) for key, value in inputs.items()}
         if self.text_encoder_model.device != device:
             self.text_encoder_model.to(device)
-
+        
         def forward_with_temperature(inputs, tau=0.01, l_min=8, l_max=10):
             handlers = []
-
+            
             def hook(module, input, output):
                 output[:] *= tau
                 return output
-
+        
             for i in range(l_min, l_max):
                 handler = self.text_encoder_model.encoder.block[i].layer[0].SelfAttention.q.register_forward_hook(hook)
                 handlers.append(handler)
-
+        
             with torch.no_grad():
                 outputs = self.text_encoder_model(**inputs)
                 last_hidden_states = outputs.last_hidden_state
-
+        
             for hook in handlers:
                 hook.remove()
-
+        
             return last_hidden_states
-
+    
         last_hidden_states = forward_with_temperature(inputs, tau, l_min, l_max)
         return last_hidden_states
 
@@ -236,7 +236,7 @@ class ACEStepPipeline:
 
     def get_lang(self, text):
         language = "en"
-        try:
+        try:    
             _ = self.lang_segment.getTexts(text)
             langCounts = self.lang_segment.getCounts()
             language = langCounts[0][0]
@@ -912,9 +912,9 @@ class ACEStepPipeline:
 
         if is_extend:
             if to_right_pad_gt_latents is not None:
-                target_latents = torch.cate([target_latents, to_right_pad_gt_latents], dim=-1)
+                target_latents = torch.cat([target_latents, to_right_pad_gt_latents], dim=-1)
             if to_left_pad_gt_latents is not None:
-                target_latents = torch.cate([to_right_pad_gt_latents, target_latents], dim=0)
+                target_latents = torch.cat([to_right_pad_gt_latents, target_latents], dim=0)
         return target_latents
 
     def latents2audio(self, latents, target_wav_duration_second=30, sample_rate=48000, save_path=None, format="flac"):