Force traditional loading: low_cpu_mem_usage=False + explicit device move

acestep/handler.py

@@ -487,23 +487,24 @@ class AceStepHandler:
                 last_attn_error = None
                 self.model = None
                 
-                # Use device context to force model initialization on CPU instead of meta device
+                # Force traditional loading to avoid meta device initialization
                 # ACE-Step's ResidualFSQ performs tensor assertions during __init__ that fail on meta device
-                with torch.device("cpu"):
-                    for candidate in attn_candidates:
-                        try:
-                            logger.info(f"[initialize_service] Attempting to load model with attention implementation: {candidate}")
-                            self.model = AutoModel.from_pretrained(
-                                acestep_v15_checkpoint_path,
-                                trust_remote_code=True,
-                                attn_implementation=candidate,
-                                dtype="bfloat16"  # Use string like official demo
-                            )
-                            attn_implementation = candidate
-                            break
-                        except Exception as e:
-                            last_attn_error = e
-                            logger.warning(f"[initialize_service] Failed to load model with {candidate}: {e}")
+                for candidate in attn_candidates:
+                    try:
+                        logger.info(f"[initialize_service] Attempting to load model with attention implementation: {candidate}")
+                        self.model = AutoModel.from_pretrained(
+                            acestep_v15_checkpoint_path,
+                            trust_remote_code=True,
+                            attn_implementation=candidate,
+                            torch_dtype=torch.bfloat16,
+                            low_cpu_mem_usage=False,  # Disable meta device
+                            device_map=None  # No automatic device mapping
+                        ).to(self.device)  # Explicitly move to target device after load
+                        attn_implementation = candidate
+                        break
+                    except Exception as e:
+                        last_attn_error = e
+                        logger.warning(f"[initialize_service] Failed to load model with {candidate}: {e}")
 
                 if self.model is None:
                     raise RuntimeError(