Update app.py

app.py

@@ -80,10 +80,14 @@ os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
 os.environ.setdefault("HF_HUB_DISABLE_TELEMETRY", "1")
 os.environ.setdefault("TRANSFORMERS_NO_ADVISORY_WARNINGS", "1")
 
+# --- Force CPU ---
 os.environ["CUDA_VISIBLE_DEVICES"] = ""
 os.environ["ACCELERATE_USE_CPU"] = "true"
 os.environ["TRANSFORMERS_ATTN_IMPLEMENTATION"] = "eager"
+
+# --- ACE-Step: обход NaN на CPU ---
 os.environ.setdefault("ACESTEP_LM_BACKEND", "pt")
+os.environ.setdefault("ACESTEP_DTYPE", "float32")
 
 _cgroup_quota = _cgroup_cpu_quota()
 print(
@@ -108,8 +112,11 @@ import spaces  # noqa: F401
 
 import torch
 
+# --- Ключевой фикс против NaN в DiT на CPU ---
+torch.set_default_dtype(torch.float32)
+
 # =========================================================================== #
-# Force CPU: перехватываем всё, что может уехать на cuda
+# Force CPU
 # =========================================================================== #
 torch.cuda.is_available = lambda: False                        # type: ignore[assignment]
 torch.cuda.device_count = lambda: 0                            # type: ignore[assignment]
@@ -205,7 +212,6 @@ def _zerogpu_startup_probe():
 # --------------------------------------
 
 
-# Патч spaces.GPU → no-op для acestep (CPU). Не влияет на probe выше.
 def _patch_spaces_gpu_noop():
     if getattr(spaces.GPU, "_acestep_noop", False):
         return
@@ -231,7 +237,7 @@ except Exception as _e:  # pragma: no cover
 
 
 # =========================================================================== #
-#                   SHARED: N-GAN (24 кГц → 48 кГц)                          #
+#                   SHARED: N-GAN
 # =========================================================================== #
 _NGAN_AVAILABLE = False
 _ngan_model = None
@@ -298,7 +304,7 @@ def enhance_to_48k(audio, sr: int):
 
 
 # =========================================================================== #
-#            SHARED: Whisper large-v3-turbo (ASR, CPU, fp32)                 #
+#            SHARED: Whisper large-v3-turbo
 # =========================================================================== #
 ASR_MODEL = "openai/whisper-large-v3-turbo"
 _asr_pipe = None
@@ -353,7 +359,7 @@ def transcribe_reference(audio_path):
 
 
 # =========================================================================== #
-#                         TAB 1: HIGGS AUDIO V3                              #
+#                         TAB 1: HIGGS AUDIO V3
 # =========================================================================== #
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
@@ -438,7 +444,7 @@ def synthesize_higgs(
 
 
 # =========================================================================== #
-#                         TAB 2: FIREREDTTS3                                 #
+#                         TAB 2: FIREREDTTS3
 # =========================================================================== #
 from transformers.modeling_utils import PreTrainedModel
 
@@ -768,7 +774,7 @@ def voice_clone_firered(
 
 
 # =========================================================================== #
-#                  TAB 3: ACE-STEP 1.5 XL-TURBO (music gen, CPU)             #
+#                  TAB 3: ACE-STEP 1.5 XL-TURBO
 # =========================================================================== #
 ACESTEP_DIT_MODEL = os.environ.get("ACESTEP_DIT_MODEL", "acestep-v15-turbo")
 ACESTEP_LM_MODEL = os.environ.get("ACESTEP_LM_MODEL", "acestep-5Hz-lm-1.7B")
@@ -781,12 +787,8 @@ _acestep_status = {"available": True, "error": None}
 
 @contextlib.contextmanager
 def _no_meta_init():
-    """Убирает meta-устройство из контекста инициализации transformers 5.x.
-
-    Иначе ResidualFSQ внутри ACE-Step вызывает .item() на мета-тензоре
-    и падает с RuntimeError('Tensor.item() cannot be called on meta tensors').
-    """
-    orig = PreTrainedModel.get_init_context.__func__  # unwrap classmethod
+    """Убирает meta-устройство из контекста инициализации transformers 5.x."""
+    orig = PreTrainedModel.get_init_context.__func__
 
     def patched(cls, *args, **kwargs):
         return [
@@ -818,7 +820,6 @@ def _get_acestep_handlers():
     if not _acestep_status["available"]:
         raise gr.Error(f"ACE-Step недоступен: {_acestep_status['error']}")
 
-    # ВАЖНО: до импорта acestep — патч spaces.GPU
     _patch_spaces_gpu_noop()
 
     try:
@@ -872,13 +873,62 @@ def _acestep_lyrics(verse: str, chorus: str, bridge: str, instrumental: bool) ->
     return "\n".join(parts) if parts else "[inst]"
 
 
+def _run_acestep_generate(params, config, save_dir):
+    """Обёртка с защитой от NaN: retry при NaN с другим shift."""
+    from acestep.inference import generate_music
+
+    dit, llm = _get_acestep_handlers()
+
+    try:
+        with torch.inference_mode():
+            result = generate_music(
+                dit_handler=dit,
+                llm_handler=llm,
+                params=params,
+                config=config,
+                save_dir=save_dir,
+            )
+    except Exception as exc:
+        raise gr.Error(f"Ошибка генерации ACE-Step: {exc}")
+
+    if not result.success:
+        err = str(getattr(result, "error", "unknown"))
+        # NaN → retry с другим shift
+        if "NaN" in err or "NaN" in err.lower():
+            print(f"[acestep] NaN detected → retry with shift=1.0", flush=True)
+            try:
+                params.shift = 1.0
+            except Exception:
+                pass
+            save_dir2 = tempfile.mkdtemp(prefix="acestep_retry_")
+            try:
+                with torch.inference_mode():
+                    result = generate_music(
+                        dit_handler=dit,
+                        llm_handler=llm,
+                        params=params,
+                        config=config,
+                        save_dir=save_dir2,
+                    )
+            except Exception as exc:
+                raise gr.Error(f"Повторная генерация тоже упала: {exc}")
+            if not result.success:
+                raise gr.Error(f"ACE-Step вернул ошибку: {result.error}")
+        else:
+            raise gr.Error(f"ACE-Step вернул ошибку: {err}")
+
+    if not result.audios:
+        raise gr.Error("ACE-Step не вернул аудио.")
+    return result.audios[0]["path"]
+
+
 def acestep_text2music(
     caption, verse, chorus, bridge, instrumental,
-    duration, steps, guidance, seed,
+    duration, steps, guidance, seed, shift,
 ):
     _patch_spaces_gpu_noop()
-    from acestep.inference import GenerationParams, GenerationConfig, generate_music
-    dit, llm = _get_acestep_handlers()
+    from acestep.inference import GenerationParams, GenerationConfig
+
     lyrics = _acestep_lyrics(verse, chorus, bridge, instrumental)
     params = GenerationParams(
         caption=caption or "",
@@ -887,6 +937,7 @@ def acestep_text2music(
         inference_steps=int(max(1, int(steps))),
         guidance_scale=float(guidance),
         seed=int(seed) if (seed is not None and int(seed) >= 0) else -1,
+        shift=float(shift) if shift is not None else 3.0,
     )
     config = GenerationConfig(
         batch_size=1,
@@ -896,25 +947,10 @@ def acestep_text2music(
     save_dir = tempfile.mkdtemp(prefix="acestep_")
     print(
         f"[acestep] gen: steps={steps} guidance={guidance} "
-        f"dur={duration}s → {save_dir}",
+        f"shift={shift} dur={duration}s → {save_dir}",
         flush=True,
     )
-    try:
-        with torch.inference_mode():
-            result = generate_music(
-                dit_handler=dit,
-                llm_handler=llm,
-                params=params,
-                config=config,
-                save_dir=save_dir,
-            )
-    except Exception as exc:
-        raise gr.Error(f"Ошибка генерации ACE-Step: {exc}")
-    if not result.success:
-        raise gr.Error(f"ACE-Step вернул ошибку: {result.error}")
-    if not result.audios:
-        raise gr.Error("ACE-Step не вернул аудио.")
-    return result.audios[0]["path"]
+    return _run_acestep_generate(params, config, save_dir)
 
 
 def acestep_create_sample(query, instrumental, vocal_language):
@@ -967,7 +1003,7 @@ def acestep_format_sample(caption, lyrics, user_bpm, user_keyscale):
 
 
 # =========================================================================== #
-#                                    UI                                       #
+#                                    UI
 # =========================================================================== #
 EN_PROMPT = os.path.join(HERE, "examples", "en_prompt.wav")
 ZH_PROMPT = os.path.join(HERE, "examples", "zh_prompt.wav")
@@ -1197,8 +1233,9 @@ with gr.Blocks(title="Higgs Audio v3 + FireRedTTS3 + ACE-Step (CPU)", fill_width
             gr.Markdown(
                 "Генерация **музыки** по описанию + куплет/припев на "
                 "**ACE-Step 1.5 XL-Turbo**, принудительно на CPU. "
-                "⚠️ Первый запуск скачивает ~18 ГБ весов (несколько минут); "
-                "сама генерация на CPU идёт медленно — начинайте с 20–30 секунд."
+                "⚠️ Первый запуск скачивает ~18 ГБ весов; генерация на CPU идёт "
+                "медленно — начинайте с 20–30 секунд. При ошибке NaN автоматически "
+                "перезапускается с `shift=1.0`."
             )
             with gr.Tabs():
                 with gr.Tab("Text → Music"):
@@ -1245,7 +1282,12 @@ with gr.Blocks(title="Higgs Audio v3 + FireRedTTS3 + ACE-Step (CPU)", fill_width
                                 )
                                 as_guidance = gr.Slider(
                                     0.0, 20.0, value=7.0, step=0.5,
-                                    label="Guidance scale",
+                                    label="Guidance scale (Turbo ignores → 1.0)",
+                                )
+                                as_shift = gr.Slider(
+                                    1.0, 5.0, value=3.0, step=0.1,
+                                    label="Shift (Timestep shift)",
+                                    info="3.0 — стабильный старт; меньше = быстрее, больше = стабильнее",
                                 )
                                 as_seed = gr.Number(
                                     value=-1, precision=0, label="Seed (-1 = random)",
@@ -1255,7 +1297,7 @@ with gr.Blocks(title="Higgs Audio v3 + FireRedTTS3 + ACE-Step (CPU)", fill_width
                         inputs=[
                             as_caption, as_verse, as_chorus, as_bridge,
                             as_instrumental, as_duration,
-                            as_steps, as_guidance, as_seed,
+                            as_steps, as_guidance, as_seed, as_shift,
                         ],
                         outputs=[as_out],
                         api_name="acestep_text2music",