Update utils.py

utils.py

@@ -132,11 +132,22 @@ def apply_tta(
     mix: torch.Tensor,
     waveforms_orig: Dict[str, torch.Tensor],
     device: str,
-    model_type: str
+    model_type: str,
+    progress=None  # Gradio progress nesnesi
 ) -> Dict[str, torch.Tensor]:
     track_proc_list = [mix[::-1].clone(), -mix.clone()]
+    total_steps = len(track_proc_list)
+    processed_steps = 0
+
     for i, augmented_mix in enumerate(track_proc_list):
-        waveforms = demix(config, model, augmented_mix, device, model_type=model_type, pbar=False)
+        # TTA adımı için ilerleme güncellemesi
+        processed_steps += 1
+        progress_value = round((processed_steps / total_steps) * 50)  # TTA için 0-50% aralığı
+        if progress is not None and callable(getattr(progress, '__call__', None)):
+            progress(progress_value / 100, desc=f"Applying TTA step {processed_steps}/{total_steps}")
+            update_progress_html(f"Applying TTA step {processed_steps}/{total_steps}", progress_value)
+
+        waveforms = demix(config, model, augmented_mix, device, model_type=model_type, pbar=False, progress=progress)
         for el in waveforms:
             if i == 0:
                 waveforms_orig[el] += waveforms[el][::-1].clone()
@@ -146,8 +157,15 @@ def apply_tta(
         gc.collect()
         if device.startswith('cuda'):
             torch.cuda.empty_cache()
+
     for el in waveforms_orig:
         waveforms_orig[el] /= (len(track_proc_list) + 1)
+
+    # TTA tamamlandı
+    if progress is not None and callable(getattr(progress, '__call__', None)):
+        progress(0.5, desc="TTA completed")
+        update_progress_html("TTA completed", 50)
+
     return waveforms_orig
 
 def _getWindowingArray(window_size: int, fade_size: int) -> torch.Tensor:
@@ -164,7 +182,8 @@ def demix(
     mix: torch.Tensor,
     device: str,
     model_type: str,
-    pbar: bool = False
+    pbar: bool = False,
+    progress=None  # Gradio progress nesnesi
 ) -> Dict[str, np.ndarray]:
     logging.info(f"Starting demix for model_type: {model_type}, chunk_size: {config.audio.chunk_size}")
     
@@ -196,6 +215,10 @@ def demix(
     model = model.to(device)
     model.eval()
 
+    # Toplam chunk sayısını hesapla
+    total_chunks = (mix.shape[1] + step - 1) // step
+    processed_chunks = 0
+
     with torch.no_grad():  # Çıkarım için gradyan yok
         with torch.cuda.amp.autocast(enabled=device.startswith('cuda'), dtype=torch.float16):
             req_shape = (num_instruments,) + mix.shape
@@ -205,7 +228,7 @@ def demix(
             i = 0
             batch_data = []
             batch_locations = []
-            progress_bar = tqdm(total=mix.shape[1], desc="Processing audio chunks", leave=False) if pbar else None
+            start_time = time.time()
 
             while i < mix.shape[1]:
                 part = mix[:, i:i + chunk_size]
@@ -240,6 +263,13 @@ def demix(
                             result[..., start:start + seg_len] += x[j, ..., :seg_len]
                             counter[..., start:start + seg_len] += 1.0
 
+                    # İlerleme güncellemesi
+                    processed_chunks += len(batch_data)
+                    progress_value = min(round((processed_chunks / total_chunks) * 100), 100)  # %1 hassasiyet
+                    if progress is not None and callable(getattr(progress, '__call__', None)):
+                        progress(progress_value / 100, desc=f"Processing chunk {processed_chunks}/{total_chunks}")
+                        update_progress_html(f"Processing chunk {processed_chunks}/{total_chunks}", progress_value)
+
                     del arr, x
                     batch_data.clear()
                     batch_locations.clear()
@@ -248,11 +278,8 @@ def demix(
                         torch.cuda.empty_cache()
                         logging.info("Cleared CUDA cache")
 
-                if progress_bar:
-                    progress_bar.update(step)
-
-            if progress_bar:
-                progress_bar.close()
+            elapsed_time = time.time() - start_time
+            logging.info(f"Demix completed in {elapsed_time:.2f} seconds")
 
             estimated_sources = result / (counter + 1e-8)
             estimated_sources = estimated_sources.numpy().astype(np.float32)
@@ -264,6 +291,12 @@ def demix(
     instruments = config.training.instruments if mode == "demucs" else prefer_target_instrument(config)
     ret_data = {k: v for k, v in zip(instruments, estimated_sources)}
     logging.info("Demix completed successfully")
+
+    # Son ilerleme güncellemesi
+    if progress is not None and callable(getattr(progress, '__call__', None)):
+        progress(1.0, desc="Demix completed")
+        update_progress_html("Demix completed", 100)
+
     return ret_data
 
 def prefer_target_instrument(config: ConfigDict) -> List[str]: