Update app.py

app.py

@@ -27,28 +27,40 @@ snac_model = None
 @spaces.GPU()
 def load_model():
     global model, tokenizer, snac_model
-    try:
-        print("Loading SNAC model...")
-        snac_model = SNAC.from_pretrained("hubertsiuzdak/snac_24khz")
-        snac_model = snac_model.to(device)
-
-        print("Loading Orpheus model...")
-        model_name = "canopylabs/orpheus-3b-0.1-ft"
-
-        snapshot_download(
-            repo_id=model_name,
-            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)
-        print(f"Orpheus model and tokenizer loaded to {device}")
-    except Exception as e:
-        print(f"Error loading model: {str(e)}")
-        raise
+    
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    
+    print("Loading SNAC model...")
+    snac_model = SNAC.from_pretrained("hubertsiuzdak/snac_24khz")
+    snac_model = snac_model.to(device)
+
+    print("Loading Orpheus model...")
+    model_name = "canopylabs/orpheus-3b-0.1-ft"
+
+    snapshot_download(
+        repo_id=model_name,
+        allow_patterns=[
+            "config.json",
+            "*.safetensors",
+            "model.safetensors.index.json",
+            "tokenizer.json",
+            "tokenizer_config.json",
+            "special_tokens_map.json",
+            "vocab.json",
+            "merges.txt",
+        ],
+        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)
+    print(f"Orpheus model and tokenizer loaded to {device}")
 
 @spaces.GPU()
 def generate_podcast_script(api_key, content, uploaded_file, duration, num_hosts):
@@ -245,7 +257,9 @@ with gr.Blocks() as demo:
 
 if __name__ == "__main__":
     try:
-        load_model()
-        demo.queue().launch()
+        print("Loading models...")
+        load_model()  # This function should be defined to load all necessary models
+        print("Models loaded successfully. Launching the interface...")
+        demo.queue().launch(share=False, ssr_mode=False)
     except Exception as e:
-        logger.error(f"Error launching the application: {str(e)}")
+        print(f"Error during startup: {e}")