Update wrapper.py

wrapper.py

@@ -4,6 +4,23 @@ import os
 import subprocess
 import sys
 
+def debug_environment():
+    """Print debug information about the environment"""
+    print("\n=== Debug Information ===")
+    print(f"Current working directory: {os.getcwd()}")
+    print("\nDirectory contents:")
+    for root, dirs, files in os.walk("."):
+        level = root.replace(".", "").count(os.sep)
+        indent = " " * 4 * level
+        print(f"{indent}{os.path.basename(root)}/")
+        subindent = " " * 4 * (level + 1)
+        for f in files:
+            print(f"{subindent}{f}")
+    
+    print("\nEnvironment variables:")
+    for key, value in os.environ.items():
+        print(f"{key}={value}")
+
 def main():
     parser = argparse.ArgumentParser(description='Run YuE model with direct input')
     parser.add_argument('--genre', type=str, required=True, help='Genre tags for the music')
@@ -15,34 +32,41 @@ def main():
     
     args = parser.parse_args()
 
+    print("\n=== Starting YuE Inference ===")
+    debug_environment()
+
     # Create temporary files for genre and lyrics
     with tempfile.NamedTemporaryFile(mode='w', delete=False) as genre_file:
         genre_file.write(args.genre)
         genre_path = genre_file.name
+        print(f"\nCreated genre file at: {genre_path}")
+        print(f"Genre content: {args.genre}")
 
     with tempfile.NamedTemporaryFile(mode='w', delete=False) as lyrics_file:
         lyrics_file.write(args.lyrics)
         lyrics_path = lyrics_file.name
+        print(f"\nCreated lyrics file at: {lyrics_path}")
+        print(f"Lyrics content: {args.lyrics}")
 
     output_dir = '/home/user/app/output'
     
     try:
-         # Add debugging before running inference
+        # Get the directory where wrapper.py is located
         current_dir = os.path.dirname(os.path.abspath(__file__))
-        print("\nDebugging information:")
-        print(f"Current directory: {current_dir}")
-        print("\nContents of current directory:")
-        print(os.listdir(current_dir))
+        infer_script = os.path.join(current_dir, 'infer.py')
         
-        if os.path.exists("./mm_tokenizer_v0.2_hf"):
-            print("\nContents of mm_tokenizer_v0.2_hf:")
-            print(os.listdir("./mm_tokenizer_v0.2_hf"))
-        else:
-            print("\nmm_tokenizer_v0.2_hf directory not found!")
+        print(f"\nInference script path: {infer_script}")
+        print(f"Script exists: {os.path.exists(infer_script)}")
         
-        # Run the inference script
-        subprocess.run([
-            'python', os.path.join(current_dir, 'infer.py'),
+        if os.path.exists('./mm_tokenizer_v0.2_hf'):
+            print("\nTokenizer directory contents:")
+            print(os.listdir('./mm_tokenizer_v0.2_hf'))
+        else:
+            print("\nWARNING: Tokenizer directory not found!")
+            
+        print("\nExecuting inference command...")
+        command = [
+            'python', infer_script,
             '--stage1_model', 'm-a-p/YuE-s1-7B-anneal-en-cot',
             '--stage2_model', 'm-a-p/YuE-s2-1B-general',
             '--genre_txt', genre_path,
@@ -52,15 +76,44 @@ def main():
             '--output_dir', output_dir,
             '--cuda_idx', str(args.cuda_idx),
             '--max_new_tokens', str(args.max_new_tokens)
-        ], check=True)
+        ]
+        
+        print(f"Command: {' '.join(command)}")
+        
+        # Run the inference script
+        result = subprocess.run(command, 
+                              check=True,
+                              capture_output=True,
+                              text=True)
+        
+        print("\nInference completed successfully!")
+        print("\nStdout:")
+        print(result.stdout)
+        
+        if result.stderr:
+            print("\nStderr:")
+            print(result.stderr)
         
         print(f"\nOutput directory: {output_dir}")
-        print("Generated files:")
-        for file in os.listdir(output_dir):
-            print(f"- {os.path.join(output_dir, file)}")
+        if os.path.exists(output_dir):
+            print("Generated files:")
+            for file in os.listdir(output_dir):
+                file_path = os.path.join(output_dir, file)
+                print(f"- {file_path} ({os.path.getsize(file_path)} bytes)")
+        else:
+            print("WARNING: Output directory does not exist!")
         
+    except subprocess.CalledProcessError as e:
+        print("\nError running inference script:")
+        print(f"Exit code: {e.returncode}")
+        print("\nStdout:")
+        print(e.stdout)
+        print("\nStderr:")
+        print(e.stderr)
+        raise
     finally:
         # Clean up temporary files
+        print("\nCleaning up temporary files...")
         os.unlink(genre_path)
         os.unlink(lyrics_path)