Update app.py

app.py

@@ -5,7 +5,7 @@ import io
 import re
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
-from huggingface_hub import snapshot_download
+from huggingface_hub import snapshot_download, login
 import torchaudio
 from torchaudio.functional import resample
 import threading
@@ -16,41 +16,47 @@ import logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Initialize Gemini AI
-genai.configure(api_key='YOUR_GEMINI_API_KEY')
-
 # Set up device
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Load Orpheus model
-print("Loading Orpheus model...")
+# Model name
 model_name = "canopylabs/orpheus-3b-0.1-ft"
 
-snapshot_download(
-    repo_id=model_name,
-    allow_patterns=[
-        "config.json",
-        "*.safetensors",
-        "model.safetensors.index.json",
-    ],
-    ignore_patterns=[
-        "optimizer.pt",
-        "pytorch_model.bin",
-        "training_args.bin",
-        "scheduler.pt",
-        "tokenizer.json",
-        "tokenizer_config.json",
-        "special_tokens_map.json",
-        "vocab.json",
-        "merges.txt",
-        "tokenizer.*"
-    ]
-)
-
-model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16)
-model.to(device)
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-print(f"Orpheus model loaded to {device}")
+def load_model(hf_token):
+    login(token=hf_token)
+    
+    print("Loading Orpheus model...")
+    snapshot_download(
+        repo_id=model_name,
+        use_auth_token=hf_token,
+        allow_patterns=[
+            "config.json",
+            "*.safetensors",
+            "model.safetensors.index.json",
+        ],
+        ignore_patterns=[
+            "optimizer.pt",
+            "pytorch_model.bin",
+            "training_args.bin",
+            "scheduler.pt",
+            "tokenizer.json",
+            "tokenizer_config.json",
+            "special_tokens_map.json",
+            "vocab.json",
+            "merges.txt",
+            "tokenizer.*"
+        ]
+    )
+
+    model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16)
+    model.to(device)
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    print(f"Orpheus model loaded to {device}")
+    return model, tokenizer
+
+# Initialize as None, will be loaded when HF token is provided
+model = None
+tokenizer = None
 
 def generate_podcast_script(api_key, content, duration, num_hosts):
     genai.configure(api_key=api_key)
@@ -90,6 +96,7 @@ def generate_podcast_script(api_key, content, duration, num_hosts):
     return clean_text
 
 def text_to_speech(text, voice):
+    global model, tokenizer
     inputs = tokenizer(text, return_tensors="pt").to(device)
     with torch.no_grad():
         output = model.generate(**inputs, max_new_tokens=256)
@@ -131,6 +138,10 @@ def render_podcast(api_key, script, voice1, voice2, num_hosts):
 with gr.Blocks() as demo:
     gr.Markdown("# AI Podcast Generator")
     
+    hf_token_input = gr.Textbox(label="Enter your Hugging Face API Token", type="password")
+    load_model_btn = gr.Button("Load Orpheus Model")
+    model_status = gr.Markdown("Model not loaded")
+    
     api_key_input = gr.Textbox(label="Enter your Gemini API Key", type="password")
     
     with gr.Row():
@@ -153,6 +164,13 @@ with gr.Blocks() as demo:
     render_btn = gr.Button("Render Podcast")
     audio_output = gr.Audio(label="Generated Podcast")
     
+    def load_model_wrapper(hf_token):
+        global model, tokenizer
+        model, tokenizer = load_model(hf_token)
+        return "Model loaded successfully"
+    
+    load_model_btn.click(load_model_wrapper, inputs=[hf_token_input], outputs=[model_status])
+    
     def generate_script_wrapper(api_key, content, duration, num_hosts):
         return generate_podcast_script(api_key, content, duration, num_hosts)