Update wrapper.py

wrapper.py

@@ -24,21 +24,31 @@ def main():
         lyrics_file.write(args.lyrics)
         lyrics_path = lyrics_file.name
 
+    output_dir = '/home/user/app/output'
+    
     try:
+        # Get the directory where wrapper.py is located
+        current_dir = os.path.dirname(os.path.abspath(__file__))
+        
         # Run the inference script
         subprocess.run([
-            'python', 'infer.py',
+            'python', os.path.join(current_dir, 'infer.py'),
             '--stage1_model', 'm-a-p/YuE-s1-7B-anneal-en-cot',
             '--stage2_model', 'm-a-p/YuE-s2-1B-general',
             '--genre_txt', genre_path,
             '--lyrics_txt', lyrics_path,
             '--run_n_segments', str(args.run_n_segments),
             '--stage2_batch_size', str(args.stage2_batch_size),
-            '--output_dir', '/app/output',
+            '--output_dir', output_dir,
             '--cuda_idx', str(args.cuda_idx),
             '--max_new_tokens', str(args.max_new_tokens)
         ], check=True)
         
+        print(f"\nOutput directory: {output_dir}")
+        print("Generated files:")
+        for file in os.listdir(output_dir):
+            print(f"- {os.path.join(output_dir, file)}")
+        
     finally:
         # Clean up temporary files
         os.unlink(genre_path)