Update app.py

app.py

@@ -16,30 +16,27 @@ import logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+def get_device():
+    if torch.cuda.is_available():
+        try:
+            torch.cuda.init()
+            return torch.device("cuda")
+        except Exception as e:
+            print(f"CUDA initialization failed: {e}")
+    return torch.device("cpu")
+
+device = get_device()
+print(f"Using device: {device}")
 
 model = None
 tokenizer = None
 
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
-from huggingface_hub import snapshot_download, login
-import os
-
-# Set up device
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-# Initialize model and tokenizer
-model = None
-tokenizer = None
-
 def load_model():
     global model, tokenizer
     
     print("Loading Orpheus model...")
     model_name = "canopylabs/orpheus-3b-0.1-ft"
 
-    # Get Hugging Face token from environment variable
     hf_token = os.environ.get("HUGGINGFACE_TOKEN")
     if not hf_token:
         raise ValueError("HUGGINGFACE_TOKEN environment variable is not set")
@@ -68,26 +65,11 @@ def load_model():
         ]
     )
 
-    model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16)
+    model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float32 if device.type == 'cpu' else torch.bfloat16)
     model.to(device)
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     print(f"Orpheus model and tokenizer loaded to {device}")
 
-# Load the model before creating the Gradio interface
-load_model()
-
-def text_to_speech(text, voice):
-    global model, tokenizer
-    if tokenizer is None or model is None:
-        raise ValueError("Model or tokenizer not initialized. Please call load_model() first.")
-    
-    inputs = tokenizer(text, return_tensors="pt").to(device)
-    with torch.no_grad():
-        output = model.generate(**inputs, max_new_tokens=256)
-    mel = output[0].cpu().numpy()
-    audio = mel_to_audio(mel)
-    return audio
-
 def generate_podcast_script(api_key, content, duration, num_hosts):
     genai.configure(api_key=api_key)
     model = genai.GenerativeModel('gemini-2.5-pro-preview-03-25')