Add TOKEN retrieval for Hugging Face API in setup function

MuseControlLite_setup.py

@@ -27,6 +27,9 @@ from utils.stable_audio_dataset_utils import load_audio_file
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
 import soundfile as sf
 
+
+TOKEN = os.environ.get("HF_TOKEN") or os.environ.get("HF_HUB_TOKEN")
+
 # For zero initialized 1D CNN in the attention processor
 def zero_module(module):
     for p in module.parameters():
@@ -507,6 +510,7 @@ class StableAudioAttnProcessor2_0_rotary_double(torch.nn.Module):
         hidden_states = hidden_states / attn.rescale_output_factor
 
         return hidden_states
+    
 def setup_MuseControlLite(config, weight_dtype, transformer_ckpt):
     """
     Setup AP-adapter pipeline with attention processors and load checkpoints.
@@ -525,7 +529,10 @@ def setup_MuseControlLite(config, weight_dtype, transformer_ckpt):
     else:
         from pipeline.stable_audio_multi_cfg_pipe import StableAudioPipeline
         attn_processor = StableAudioAttnProcessor2_0_rotary
-    pipe = StableAudioPipeline.from_pretrained("stabilityai/stable-audio-open-1.0", torch_dtype=weight_dtype)
+    pipe = StableAudioPipeline.from_pretrained(
+        "stabilityai/stable-audio-open-1.0", 
+        torch_dtype=weight_dtype, token=TOKEN
+        )
     pipe.scheduler.config.sigma_max = config["sigma_max"]
     pipe.scheduler.config.sigma_min = config["sigma_min"]
     transformer = pipe.transformer