Update app.py

app.py

@@ -97,10 +97,13 @@ model = AutoModelForCausalLM.from_pretrained(
     attn_implementation="flash_attention_2",
 ).to(device)
 model.eval()
-try:
-    model = torch.compile(model)
-except Exception as e:
-    print("torch.compile skipped for Stage1 model:", e)
+
+model_stage2 = AutoModelForCausalLM.from_pretrained(
+    STAGE2_MODEL,
+    torch_dtype=torch.float16,
+    attn_implementation="flash_attention_2",
+).to(device)
+model_stage2.eval()
 
 mmtokenizer = _MMSentencePieceTokenizer("./mm_tokenizer_v0.2_hf/tokenizer.model")
 
@@ -115,10 +118,6 @@ codec_model = codec_class(**model_config.generator.config).to(device)
 parameter_dict = torch.load(RESUME_PATH, map_location="cpu")
 codec_model.load_state_dict(parameter_dict["codec_model"])
 codec_model.eval()
-try:
-    codec_model = torch.compile(codec_model)
-except Exception as e:
-    print("torch.compile skipped for codec_model:", e)
 
 # Precompile regex for splitting lyrics
 LYRICS_PATTERN = re.compile(r"\[(\w+)\](.*?)\n(?=\[|\Z)", re.DOTALL)
@@ -409,12 +408,7 @@ def generate_music(
 
         # ---------------- Stage 2: Refinement/Upsampling ----------------
         print("Stage 2 inference...")
-        model_stage2 = AutoModelForCausalLM.from_pretrained(
-            STAGE2_MODEL,
-            torch_dtype=torch.float16,
-            attn_implementation="flash_attention_2",
-        ).to(device)
-        model_stage2.eval()
+        
         stage2_result = stage2_inference(model_stage2, stage1_output_set, stage2_output_dir, batch_size=STAGE2_BATCH_SIZE)
         print("Stage 2 inference completed.")