Update app.py

app.py

@@ -31,39 +31,26 @@ snac_model = None
 def load_model():
     global model, tokenizer, snac_model
     try:
-        logger.info("Loading SNAC model...")
+        print("Loading SNAC model...")
         snac_model = SNAC.from_pretrained("hubertsiuzdak/snac_24khz")
         snac_model = snac_model.to(device)
-        
-        if snac_model is None:
-            raise ValueError("Failed to load SNAC model")
-        
-        logger.info("Loading Orpheus model...")
-        model_name = "canopylabs/orpheus-3b-0.1-ft"
 
-        hf_token = os.environ.get("HUGGINGFACE_TOKEN")
-        if not hf_token:
-            raise ValueError("HUGGINGFACE_TOKEN environment variable is not set")
-
-        login(token=hf_token)
+        print("Loading Orpheus model...")
+        model_name = "canopylabs/orpheus-3b-0.1-ft"
 
         snapshot_download(
             repo_id=model_name,
-            use_auth_token=hf_token,
+            use_auth_token=os.environ.get("HUGGINGFACE_TOKEN"),
             allow_patterns=["config.json", "*.safetensors", "model.safetensors.index.json", "vocab.json", "merges.txt", "tokenizer.json"],
             ignore_patterns=["optimizer.pt", "pytorch_model.bin", "training_args.bin", "scheduler.pt"]
         )
 
         model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16)
         model.to(device)
-        
         tokenizer = AutoTokenizer.from_pretrained(model_name)
-        if tokenizer is None:
-            raise ValueError("Failed to load tokenizer")
-        
-        logger.info(f"Orpheus model and tokenizer loaded to {device}")
+        print(f"Orpheus model and tokenizer loaded to {device}")
     except Exception as e:
-        logger.error(f"Error loading model: {str(e)}")
+        print(f"Error loading model: {str(e)}")
         raise
 
 @spaces.GPU()
@@ -167,6 +154,11 @@ def redistribute_codes(code_list, snac_model):
 
 @spaces.GPU()
 def generate_speech(text, voice, temperature, top_p, repetition_penalty, max_new_tokens):
+    global tokenizer
+    if tokenizer is None:
+        print("Tokenizer is not initialized. Please ensure the model is properly loaded.")
+        return None
+
     if not text.strip():
         return None
     
@@ -191,7 +183,7 @@ def generate_speech(text, voice, temperature, top_p, repetition_penalty, max_new
         
         return (24000, audio_samples)  # Return sample rate and audio
     except Exception as e:
-        logger.error(f"Error generating speech: {e}", exc_info=True)
+        print(f"Error generating speech: {e}")
         return None
 
 @spaces.GPU()