Update app.py

app.py

@@ -103,14 +103,17 @@ def load_new_model(model_id):
 def run_inference(user_prompt, max_tokens, temperature, top_k, top_p, rep_penalty, ngram_size, do_sample, gpu):
     if gpu:
         yield from run_inference_gpu(user_prompt, max_tokens, temperature, top_k, top_p, rep_penalty, ngram_size, do_sample)
-    yield from run_inference_raw(user_prompt, max_tokens, temperature, top_k, top_p, rep_penalty, ngram_size, do_sample)
+    else:
+        yield from run_inference_raw(user_prompt, max_tokens, temperature, top_k, top_p, rep_penalty, ngram_size, do_sample)
 
 
 @spaces.GPU
 def run_inference_gpu(user_prompt, max_tokens, temperature, top_k, top_p, rep_penalty, ngram_size, do_sample):
-    yield from run_inference_raw(user_prompt, max_tokens, temperature, top_k, top_p, rep_penalty, ngram_size, do_sample)
+    if model_manager.model is not None:
+        model_manager.model = model_manager.model.to("cuda")
+    yield from run_inference_raw(user_prompt, max_tokens, temperature, top_k, top_p, rep_penalty, ngram_size, do_sample, use_cuda=True)
 
-def run_inference_raw(user_prompt, max_tokens, temperature, top_k, top_p, rep_penalty, ngram_size, do_sample):
+def run_inference_raw(user_prompt, max_tokens, temperature, top_k, top_p, rep_penalty, ngram_size, do_sample, use_cuda=False):
     """Generates text via streaming generator."""
     if model_manager.model is None or model_manager.tokenizer is None:
         yield "Please load a model first.", "Model not loaded"
@@ -140,6 +143,12 @@ def run_inference_raw(user_prompt, max_tokens, temperature, top_k, top_p, rep_pe
     # Tokenize input
     inputs = tokenizer([prompt_to_encode], return_tensors="pt")
     
+    if use_cuda:
+        inputs = {k: v.to("cuda") for k, v in inputs.items()}
+    else:
+        model = model.to("cpu")
+        inputs = {k: v.to("cpu") for k, v in inputs.items()}
+    
     # Set up the streamer
     streamer = TextIteratorStreamer(tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=skip_special)
     
@@ -199,7 +208,8 @@ def run_inference_raw(user_prompt, max_tokens, temperature, top_k, top_p, rep_pe
             display_text = display_text.replace("<|begin_of_solution|>", "✅ Final Answer:\n\n")
             display_text = display_text.replace("<|end_of_solution|>", "")
 
-        yield base_display + display_text, f"Speed: {tps:.2f} tokens/sec"
+        device_label = "CUDA" if use_cuda else "CPU"
+        yield base_display + display_text, f"Speed: {tps:.2f} tokens/sec ({device_label})"
 
 def clean_cache():
     if os.path.exists(HF_CACHE_DIR):