Update app.py

app.py

@@ -102,13 +102,13 @@ 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:
-        return run_inference_gpu(user_prompt, max_tokens, temperature, top_k, top_p, rep_penalty, ngram_size, do_sample)
-    return run_inference_raw(user_prompt, max_tokens, temperature, top_k, top_p, rep_penalty, ngram_size, do_sample)
+        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)
 
 
 @spaces.GPU
 def run_inference_gpu(user_prompt, max_tokens, temperature, top_k, top_p, rep_penalty, ngram_size, do_sample):
-    return run_inference_raw(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)
 
 def run_inference_raw(user_prompt, max_tokens, temperature, top_k, top_p, rep_penalty, ngram_size, do_sample):
     """Generates text via streaming generator."""
@@ -226,7 +226,7 @@ with gr.Blocks(title="Small MF Model Tester", theme=gr.themes.Soft()) as app:
                 with gr.Row():
                     model_id_input = gr.Dropdown(choices=MODELS, label="Model", allow_custom_value=True, show_label=False, scale=3)
                     load_btn = gr.Button("Load", variant="secondary", scale=1)
-                use_gpu = gr.Checkbox("Use GPU?")
+                use_gpu = gr.Checkbox(label="Use GPU?", value=False)
                 clean_btn = gr.Button("Clean HF Cache", variant="stop", size="sm")
 
             with gr.Accordion("Generation Configuration", open=False):