Update acestep/inference.py

acestep/inference.py

@@ -11,10 +11,8 @@ from loguru import logger
 import torch
 from acestep.audio_utils import AudioSaver, generate_uuid_from_params
 
-# HuggingFace Space environment detection
 IS_HUGGINGFACE_SPACE = os.environ.get("SPACE_ID") is not None
 
-# تعریف کلاس اختصاصی برای خروج امن قبل از اتمام محدودیت ۱۲۰ ثانیه‌ای
 class TimeoutSavedException(Exception):
     def __init__(self, task_id: str, step: int):
         self.task_id = task_id
@@ -23,31 +21,19 @@ class TimeoutSavedException(Exception):
 
 @dataclass
 class GenerationParams:
-    """Configuration for music generation parameters."""
-    # Required Inputs
     task_type: str = "text2music"
     instruction: str = "Fill the audio semantic mask based on the given conditions:"
-
-    # Audio Uploads
     reference_audio: Optional[str] = None
     src_audio: Optional[str] = None
-
-    # LM Codes Hints
     audio_codes: str = ""
-
-    # Text Inputs
     caption: str = ""
     lyrics: str = ""
     instrumental: bool = False
-
-    # Metadata
     vocal_language: str = "unknown"
     bpm: Optional[int] = None
     keyscale: str = ""
     timesignature: str = ""
     duration: float = -1.0
-
-    # Advanced Settings
     inference_steps: int = 8
     seed: int = -1
     guidance_scale: float = 7.0
@@ -55,14 +41,11 @@ class GenerationParams:
     cfg_interval_start: float = 0.0
     cfg_interval_end: float = 1.0
     shift: float = 1.0
-    infer_method: str = "ode"  # "ode" or "sde" - diffusion inference method
+    infer_method: str = "ode"
     timesteps: Optional[List[float]] = None
-
     repainting_start: float = 0.0
     repainting_end: float = -1
     audio_cover_strength: float = 1.0
-
-    # 5Hz Language Model Parameters
     thinking: bool = True
     lm_temperature: float = 0.85
     lm_cfg_scale: float = 2.0
@@ -71,10 +54,9 @@ class GenerationParams:
     lm_negative_prompt: str = "NO USER INPUT"
     use_cot_metas: bool = True
     use_cot_caption: bool = True
-    use_cot_lyrics: bool = False  # TODO: not used yet
+    use_cot_lyrics: bool = False
     use_cot_language: bool = True
     use_constrained_decoding: bool = True
-
     cot_bpm: Optional[int] = None
     cot_keyscale: str = ""
     cot_timesignature: str = ""
@@ -82,34 +64,26 @@ class GenerationParams:
     cot_vocal_language: str = "unknown"
     cot_caption: str = ""
     cot_lyrics: str = ""
-    
-    # شناسه ذخیره و احیا
     task_id: Optional[str] = None
 
     def to_dict(self) -> Dict[str, Any]:
-        """Convert config to dictionary for JSON serialization."""
         return asdict(self)
 
-
 @dataclass
 class GenerationConfig:
-    """Configuration for music generation."""
     batch_size: int = 2
     allow_lm_batch: bool = False
     use_random_seed: bool = True
     seeds: Optional[List[int]] = None
     lm_batch_chunk_size: int = 8
     constrained_decoding_debug: bool = False
-    audio_format: str = "flac"  # Default to FLAC for fast saving
+    audio_format: str = "flac"
 
     def to_dict(self) -> Dict[str, Any]:
-        """Convert config to dictionary for JSON serialization."""
         return asdict(self)
 
-
 @dataclass
 class GenerationResult:
-    """Result of music generation."""
     audios: List[Dict[str, Any]] = field(default_factory=list)
     status_message: str = ""
     extra_outputs: Dict[str, Any] = field(default_factory=dict)
@@ -117,13 +91,10 @@ class GenerationResult:
     error: Optional[str] = None
 
     def to_dict(self) -> Dict[str, Any]:
-        """Convert result to dictionary for JSON serialization."""
         return asdict(self)
 
-
 @dataclass
 class UnderstandResult:
-    """Result of music understanding from audio codes."""
     caption: str = ""
     lyrics: str = ""
     bpm: Optional[int] = None
@@ -136,13 +107,10 @@ class UnderstandResult:
     error: Optional[str] = None
 
     def to_dict(self) -> Dict[str, Any]:
-        """Convert result to dictionary for JSON serialization."""
         return asdict(self)
 
-
 @dataclass
 class CreateSampleResult:
-    """Result of creating a music sample from a natural language query."""
     caption: str = ""
     lyrics: str = ""
     bpm: Optional[int] = None
@@ -156,13 +124,10 @@ class CreateSampleResult:
     error: Optional[str] = None
 
     def to_dict(self) -> Dict[str, Any]:
-        """Convert result to dictionary for JSON serialization."""
         return asdict(self)
 
-
 @dataclass
 class FormatSampleResult:
-    """Result of formatting user-provided caption and lyrics."""
     caption: str = ""
     lyrics: str = ""
     bpm: Optional[int] = None
@@ -175,10 +140,8 @@ class FormatSampleResult:
     error: Optional[str] = None
 
     def to_dict(self) -> Dict[str, Any]:
-        """Convert result to dictionary for JSON serialization."""
         return asdict(self)
 
-
 def _update_metadata_from_lm(
     metadata: Dict[str, Any],
     bpm: Optional[int],
@@ -189,7 +152,6 @@ def _update_metadata_from_lm(
     caption: str,
     lyrics: str,
 ) -> Tuple[Optional[int], str, str, Optional[float], str, str, str]:
-    """Update metadata fields from LM output if not provided by user."""
     if bpm is None and metadata.get('bpm'):
         bpm_value = metadata.get('bpm')
         if bpm_value not in ["N/A", ""]:
@@ -224,7 +186,6 @@ def _update_metadata_from_lm(
         lyrics = metadata.get('lyrics')
     return bpm, key_scale, time_signature, audio_duration, vocal_language, caption, lyrics
 
-
 def generate_music(
     dit_handler,
     llm_handler,
@@ -233,7 +194,6 @@ def generate_music(
     save_dir: Optional[str] = None,
     progress=None,
 ) -> GenerationResult:
-    """Generate music using ACE-Step model with optional LM reasoning."""
     try:
         audio_code_string_to_use = params.audio_codes
         lm_generated_metadata = None
@@ -244,7 +204,6 @@ def generate_music(
             "total_time": 0.0,
         }
 
-        # Extract mutable copies of metadata
         bpm = params.bpm
         key_scale = params.keyscale
         time_signature = params.timesignature
@@ -257,42 +216,37 @@ def generate_music(
         need_audio_codes = not user_provided_audio_codes
         actual_batch_size = config.batch_size if config.batch_size is not None else 1
 
-        # تعیین روت دایرکتوری پایدار برای ذخیره لیتنت‌ها
         storage_root = "/data" if os.path.exists("/data") else tempfile.gettempdir()
         resume_dir = os.path.join(storage_root, "resume_tasks")
         os.makedirs(resume_dir, exist_ok=True)
 
-        # بررسی هوشمند فیلد Seed برای استخراج شناسه احیا (Tracking ID)
         active_task_id = params.task_id
         is_resumed_run = False
         
-        # اگر کاربر شناسه احیا را به جای Seed وارد کرده باشد (غیر عددی بودن سید)
         if isinstance(params.seed, str) and params.seed.strip():
             cleaned_seed = params.seed.strip()
             if not cleaned_seed.replace("-", "").replace(",", "").isdigit():
                 active_task_id = cleaned_seed
                 is_resumed_run = True
-                logger.info(f"🟢 [Save & Resume] Resuming request from Seed-Input Tracking ID: {active_task_id}")
+                logger.info(f"Resuming request from Seed-Input Tracking ID: {active_task_id}")
         
         if not active_task_id or not active_task_id.strip():
             active_task_id = str(uuid.uuid4())[:8]
             
         params.task_id = active_task_id
 
-        # بررسی فیزیکی وجود وضعیت ذخیره شده بر روی دیسک
         resume_path = os.path.join(resume_dir, f"{active_task_id}_latents.pt")
         meta_path = os.path.join(resume_dir, f"{active_task_id}_meta.json")
         
         if os.path.exists(resume_path) and os.path.exists(meta_path):
             is_resumed_run = True
-            logger.info(f"🟢 [Save & Resume] Valid saved state files found for Task {active_task_id}")
+            logger.info(f"State files found for Task {active_task_id}")
 
         seed_for_generation = ""
         if config.seeds is not None and len(config.seeds) > 0:
             if isinstance(config.seeds, list):
                 seed_for_generation = ",".join(str(s) for s in config.seeds)
 
-        # اگر در حالت احیا هستیم، از ساخت سیدهای تصادفی جدید عبور می‌کنیم تا شرایط ثابت بماند
         if is_resumed_run:
             config.use_random_seed = False
             if os.path.exists(meta_path):
@@ -308,7 +262,6 @@ def generate_music(
 
         actual_seed_list, _ = dit_handler.prepare_seeds(actual_batch_size, seed_for_generation, config.use_random_seed)
 
-        # ذخیره سیدها در ابتدای فاز یک برای بازیابی‌های احتمالی بعدی در فاز دو
         if not is_resumed_run:
             temp_meta_save = os.path.join(resume_dir, f"{active_task_id}_meta.json")
             try:
@@ -317,7 +270,6 @@ def generate_music(
             except Exception as e:
                 logger.error(f"Failed to pre-save meta seeds: {e}")
 
-        # بارگذاری مجدد دیتای متادیتا در فاز یک در صورت احیا شدن
         if is_resumed_run and os.path.exists(meta_path):
             try:
                 with open(meta_path, 'r', encoding='utf-8') as f:
@@ -325,7 +277,7 @@ def generate_music(
                 saved_lm_metadata = meta_data.get("lm_metadata")
                 if saved_lm_metadata:
                     lm_generated_metadata = saved_lm_metadata
-                    logger.info(f"🟢 [Save & Resume] Restored CoT Metadata for Task {active_task_id}")
+                    logger.info(f"Restored CoT Metadata for Task {active_task_id}")
             except Exception as e:
                 logger.error(f"Failed to restore metadata from meta file: {e}")
 
@@ -473,7 +425,6 @@ def generate_music(
             if params.use_cot_language:
                 dit_input_vocal_language = lm_generated_metadata.get("vocal_language", dit_input_vocal_language)
 
-            # ذخیره کردن وضعیت متا برای ب��زیابی در فاز دو
             temp_meta_save = os.path.join(resume_dir, f"{active_task_id}_meta.json")
             try:
                 with open(temp_meta_save, 'w', encoding='utf-8') as f:
@@ -486,14 +437,13 @@ def generate_music(
             except Exception as e:
                 logger.error(f"Failed to save temporary CoT meta: {e}")
 
-        # دزدیدن متد هندلر دیفیوشن برای رصد زمان و لیتنت‌ها
+        # اعمال دکوراتور داینامیک بر روی کلاس مدل
         decoder = dit_handler.model.decoder
         original_forward = decoder.forward
         current_step = 0
-        timeout_limit = 95.0 # ثانیه - محدودیت زمانی امن جهت خروج قبل از اتمام تایم اوت هاگینگ فیس
+        timeout_limit = 95.0
         fsm_start_time = time.time()
         
-        # بازیابی لیتنت‌های ذخیره شده قبلی و همگام‌سازی گام شروع
         is_resumed_run_active = False
         resume_step_idx = 0
         saved_latents_tensor = None
@@ -505,35 +455,29 @@ def generate_music(
                 resume_step_idx = meta_data.get("resume_step", 0)
                 saved_latents_tensor = torch.load(resume_path, map_location='cpu')
                 is_resumed_run_active = True
-                logger.info(f"🟢 [Dynamic Monkeypatch] State active. Will inject saved VRAM at step {resume_step_idx}")
+                logger.info(f"Active. Will inject saved VRAM at step {resume_step_idx}")
             except Exception as e:
-                logger.error(f"Monkeypatch state restore failed: {e}")
+                logger.error(f"Restore failed: {e}")
                 is_resumed_run_active = False
 
         def patched_forward(self_decoder, hidden_states, timestep, *args, **kwargs):
             nonlocal current_step, fsm_start_time, is_resumed_run_active, resume_step_idx, saved_latents_tensor, active_task_id
             
-            # اگر در مرحله احیا هستیم و هنوز به گام ذخیره شده نرسیده‌ایم، فورا از پردازش گام عبور کن
             if is_resumed_run_active and current_step < resume_step_idx:
                 current_step += 1
                 return (torch.zeros_like(hidden_states),)
             
-            # اگر دقیقاً به گام قطع شده رسیدیم، لیتنت‌های قبلی را به VRAM تزریق کن
             if is_resumed_run_active and current_step == resume_step_idx:
                 saved_latents_tensor = saved_latents_tensor.to(device=hidden_states.device, dtype=hidden_states.dtype)
                 hidden_states.copy_(saved_latents_tensor)
                 is_resumed_run_active = False
-                logger.info(f"🟢 [Dynamic Monkeypatch] Succeeded in overriding hidden_states at step {current_step}")
+                logger.info(f"Successfully injected saved VRAM at step {current_step}")
             
-            # محاسبه زمان باقیمانده برای فرار از تایم اوت
             elapsed_time = time.time() - fsm_start_time
             if elapsed_time > timeout_limit:
-                logger.warning(f"⚠️ [ZeroGPU Timeout Saved] Intercepted at {elapsed_time:.1f}s. Saving state for {active_task_id} at step {current_step}")
-                
-                # انتقال لیتنت‌ها به CPU و ذخیره روی هارد پایدار
+                logger.warning(f"Intercepted at {elapsed_time:.1f}s. Saving state for {active_task_id} at step {current_step}")
                 torch.save(hidden_states.cpu(), resume_path)
                 
-                # بروزرسانی فایل متا
                 current_meta = {"task_id": active_task_id, "resume_step": current_step, "seeds": actual_seed_list}
                 if lm_generated_metadata:
                     current_meta["lm_metadata"] = lm_generated_metadata
@@ -541,7 +485,6 @@ def generate_music(
                 with open(meta_path, 'w', encoding='utf-8') as f:
                     json.dump(current_meta, f, ensure_ascii=False)
                 
-                # بازیابی فوری متد اصلی دکودر
                 decoder.forward = original_forward
                 raise TimeoutSavedException(active_task_id, current_step)
             
@@ -585,7 +528,6 @@ def generate_music(
         finally:
             decoder.forward = original_forward
 
-        # در صورت اتمام موفقیت‌آمیز رندر، فایل‌های ذخیره وضعیت میانی را برای پاکیزگی دیسک حذف کن
         if os.path.exists(resume_path):
             try:
                 os.remove(resume_path)
@@ -690,7 +632,6 @@ def generate_music(
         )
 
     except TimeoutSavedException as e:
-        # مدیریت استثنای قطع امن و پرتاب آن به سمت لایه بالاتر
         raise e
     except Exception as e:
         logger.exception("Music generation failed")