Update app.py

app.py

@@ -16,45 +16,48 @@ 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 = "canopylabs/orpheus-3b-0.1-ft"
-
-HF_TOKEN = "YOUR_HUGGINGFACE_TOKEN"
-login(token=HF_TOKEN)
-
-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}")
+# Initialize model and tokenizer as None
+model = None
+tokenizer = None
+
+def load_model(hf_token):
+    global model, tokenizer
+    
+    print("Loading Orpheus model...")
+    model_name = "canopylabs/orpheus-3b-0.1-ft"
+
+    login(token=hf_token)
+
+    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 loaded successfully"
 
 def generate_podcast_script(api_key, content, duration, num_hosts):
     genai.configure(api_key=api_key)
@@ -94,6 +97,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)
@@ -135,6 +139,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():
@@ -159,6 +167,8 @@ with gr.Blocks() as demo:
     render_btn = gr.Button("Render Podcast")
     audio_output = gr.Audio(label="Generated Podcast")
     
+    load_model_btn.click(load_model, 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)