Update app.py

app.py

@@ -6,7 +6,6 @@ import shutil
 from pathlib import Path
 import zipfile
 import json
-import time
 
 class RealRVCTrainer:
     def __init__(self):
@@ -116,7 +115,7 @@ class RealRVCTrainer:
             
             progress(0.3, desc="Preprocessing audio...")
             cmd = [sys.executable, str(preprocess_script), str(dataset_path), str(sample_rate), "2"]
-            result = subprocess.run(cmd, capture_output=True, text=True, timeout=600)
+            result = subprocess.run(cmd, capture_output=True, text=True)
             
             progress(1.0, desc="Preprocessing complete!")
             return f"✅ Preprocessing Complete!\n\n🎵 Sample Rate: {sample_rate}Hz\n🔍 Features extracted\n📊 Ready for training!"
@@ -131,148 +130,83 @@ class RealRVCTrainer:
             log_dir = self.rvc_dir / "logs" / model_name
             log_dir.mkdir(parents=True, exist_ok=True)
             
-            # Create weights directory explicitly
-            weights_dir = log_dir / "weights"
-            weights_dir.mkdir(exist_ok=True)
-            
             progress(0.1, desc="Starting RVC training...")
             train_script = self.rvc_dir / "infer" / "modules" / "train" / "train.py"
             if not train_script.exists():
                 train_script = self.rvc_dir / "train_nsf_sim_cache_sid_load_pretrain.py"
             
-            if not train_script.exists():
-                return f"❌ Training script not found. Searched:\n- {self.rvc_dir / 'infer' / 'modules' / 'train' / 'train.py'}\n- {self.rvc_dir / 'train_nsf_sim_cache_sid_load_pretrain.py'}"
-            
             cmd = [
                 sys.executable, str(train_script),
-                "-e", model_name, 
-                "-sr", str(sample_rate),
-                "-f0", "1", 
-                "-bs", str(batch_size),
-                "-g", "0", 
-                "-te", str(epochs), 
-                "-se", "10",
+                "-e", model_name, "-sr", str(sample_rate),
+                "-f0", "1", "-bs", str(batch_size),
+                "-g", "0", "-te", str(epochs), "-se", "10",
                 "-pg", str(self.rvc_dir / "pretrained" / "f0G40k.pth"),
                 "-pd", str(self.rvc_dir / "pretrained" / "f0D40k.pth"),
-                "-l", "0", 
-                "-c", "0"
+                "-l", "0", "-c", "0"
             ]
             
             progress(0.2, desc=f"Training {model_name}...")
             
-            # Set up environment to include RVC directory in Python path
-            env = os.environ.copy()
-            env['PYTHONPATH'] = str(self.rvc_dir) + (f":{env.get('PYTHONPATH', '')}" if env.get('PYTHONPATH') else "")
-            
-            # Run training from RVC directory and capture output
-            process = subprocess.Popen(
-                cmd, 
-                stdout=subprocess.PIPE, 
-                stderr=subprocess.STDOUT,
-                text=True,
-                bufsize=1,
-                cwd=str(self.rvc_dir),
-                env=env
-            )
-            
-            last_output = []
-            while True:
-                line = process.stdout.readline()
-                if not line and process.poll() is not None:
-                    break
-                if line:
-                    last_output.append(line.strip())
-                    if len(last_output) > 10:
-                        last_output.pop(0)
-                    if "epoch" in line.lower() or "step" in line.lower():
-                        progress(0.2, desc=f"Training: {line.strip()[:60]}")
-            
-            return_code = process.wait()
+            # Capture both stdout and stderr
+            process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
             
-            if return_code != 0:
-                return f"❌ Training failed with code {return_code}\n\nLast output:\n" + "\n".join(last_output[-5:])
+            training_output = []
+            for line in process.stdout:
+                training_output.append(line.strip())
+                if "epoch" in line.lower():
+                    progress(0.2 + 0.6, desc=f"Training: {line.strip()[:50]}")
             
-            # Wait a moment for file system to sync
-            time.sleep(2)
+            return_code = process.wait()
             
             progress(0.9, desc="Searching for model files...")
             
-            # Expanded search paths
-            possible_paths = [
-                weights_dir,
-                log_dir,
-                log_dir / "G_0.pth",  # Sometimes saved directly
-                log_dir / "D_0.pth",
-                self.rvc_dir / "weights" / model_name,
-                self.rvc_dir / "logs" / model_name,
-                self.rvc_dir / "checkpoints" / model_name,
-            ]
+            # Search for G_*.pth (generator checkpoints) and D_*.pth (discriminator)
+            g_files = list(log_dir.glob("G_*.pth"))
+            d_files = list(log_dir.glob("D_*.pth"))
+            index_files = list(log_dir.rglob("*.index"))
+            
+            # Also check in added_* subfolders
+            added_folders = list(log_dir.glob("added_*"))
+            for folder in added_folders:
+                g_files.extend(list(folder.glob("G_*.pth")))
+                d_files.extend(list(folder.glob("D_*.pth")))
+                index_files.extend(list(folder.glob("*.index")))
             
-            # Search for model files more thoroughly
-            model_files = []
-            index_files = []
-            
-            for path in possible_paths:
-                if path.exists():
-                    if path.is_dir():
-                        model_files.extend(list(path.glob("**/*.pth")))
-                        index_files.extend(list(path.glob("**/*.index")))
-                    elif path.suffix == '.pth':
-                        model_files.append(path)
-            
-            # Search for G_*.pth and D_*.pth checkpoint files
-            if log_dir.exists():
-                g_checkpoints = list(log_dir.glob("G_*.pth"))
-                d_checkpoints = list(log_dir.glob("D_*.pth"))
-                model_files.extend(g_checkpoints)
-                model_files.extend(list(log_dir.rglob("*.index")))
-                index_files.extend(list(log_dir.rglob("*.index")))
-            
-            # Remove duplicates
-            model_files = list(set(model_files))
-            index_files = list(set(index_files))
-            
-            if model_files or index_files:
+            if g_files or index_files:
                 output_dir = self.workspace / model_name
                 output_dir.mkdir(exist_ok=True)
                 
                 files_info = []
                 
-                # Find the latest G checkpoint (Generator - the main model)
-                g_files = [f for f in model_files if f.name.startswith('G_')]
+                # Get the latest G file (highest epoch number)
                 if g_files:
-                    latest_g = max(g_files, key=lambda p: int(p.stem.split('_')[1]) if p.stem.split('_')[1].isdigit() else 0)
+                    latest_g = max(g_files, key=lambda f: int(f.stem.split('_')[1]) if f.stem.split('_')[1].isdigit() else 0)
                     shutil.copy2(latest_g, output_dir / f"{model_name}.pth")
                     model_size = latest_g.stat().st_size / (1024*1024)
-                    files_info.append(f"- {model_name}.pth (G checkpoint, {model_size:.1f}MB)")
+                    files_info.append(f"- {model_name}.pth ({model_size:.1f}MB) [from {latest_g.name}]")
                 
-                # Copy index files
-                for index_file in index_files:
-                    shutil.copy2(index_file, output_dir / index_file.name)
-                    files_info.append(f"- {index_file.name}")
+                if index_files:
+                    latest_index = max(index_files, key=lambda p: p.stat().st_mtime)
+                    shutil.copy2(latest_index, output_dir / latest_index.name)
+                    files_info.append(f"- {latest_index.name}")
                 
                 progress(1.0, desc="Training complete!")
-                return f"✅ Training Complete!\n\n🎓 Model: {model_name}\n📊 Epochs: {epochs}\n\n💾 Extracted Files:\n{chr(10).join(files_info)}\n\n📂 Location: {output_dir}\n\n🎉 Ready to download!\n\n💡 Note: G_*.pth files are RVC checkpoints. Extracted as {model_name}.pth for inference."
+                return f"✅ Training Complete!\n\n🎓 Model: {model_name}\n📊 Epochs: {epochs}\n\n💾 Model Files:\n{chr(10).join(files_info)}\n\n📂 Location: {output_dir}\n\n🎉 Ready to download!"
             else:
-                # Enhanced debug info
-                debug_info = [f"Searched paths:"]
-                for p in possible_paths:
-                    debug_info.append(f"  - {p} (exists: {p.exists()})")
-                
-                debug_info.append(f"\nLog directory contents:")
+                debug_info = [f"Return code: {return_code}", ""]
+                debug_info.append("Files in log directory:")
                 if log_dir.exists():
                     for item in log_dir.rglob("*"):
-                        size = item.stat().st_size if item.is_file() else 0
-                        debug_info.append(f"  - {item.relative_to(log_dir)} ({size} bytes)")
+                        if item.is_file():
+                            debug_info.append(f"  - {item.relative_to(log_dir)} ({item.stat().st_size} bytes)")
                 
-                debug_info.append(f"\nTraining output (last 10 lines):")
-                debug_info.extend(last_output[-10:])
+                debug_info.append("\nLast 10 lines of training output:")
+                debug_info.extend(training_output[-10:])
                 
-                return f"⚠️ Training completed but model files not found.\n\n{chr(10).join(debug_info)}\n\n💡 Tips:\n- Check if training actually ran (see output above)\n- Verify pretrained models exist\n- Try running with fewer epochs first (10-20)\n- Check disk space"
+                return f"⚠️ Training completed but model files not found.\n\n📋 Debug:\n{chr(10).join(debug_info)}\n\n💡 Check if training actually ran or failed silently."
                 
         except Exception as e:
-            import traceback
-            return f"❌ Training failed: {str(e)}\n\nTraceback:\n{traceback.format_exc()}"
+            return f"❌ Training failed: {str(e)}"
     
     def package_model(self, model_name):
         """Package model for download"""
@@ -281,8 +215,8 @@ class RealRVCTrainer:
             if not output_dir.exists():
                 output_dir = self.rvc_dir / "logs" / model_name / "weights"
             
-            if not output_dir.exists() or not any(output_dir.iterdir()):
-                return None, "❌ Model not found or empty"
+            if not output_dir.exists():
+                return None, "❌ Model not found"
             
             zip_path = self.workspace / f"{model_name}_RVC.zip"
             with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
@@ -290,8 +224,7 @@ class RealRVCTrainer:
                     if file.is_file() and (file.suffix in ['.pth', '.index', '.json']):
                         zipf.write(file, file.name)
             
-            zip_size = zip_path.stat().st_size / (1024*1024)
-            return str(zip_path), f"✅ Model packaged: {zip_path.name} ({zip_size:.1f}MB)"
+            return str(zip_path), f"✅ Model packaged: {zip_path.name}"
         except Exception as e:
             return None, f"❌ Error: {str(e)}"