Update app.py

app.py

@@ -176,9 +176,9 @@ CPUs will be slower but there is no time limit.
 
     def model_fn(prefix, history, server):
         if server == 'Huggingface ZeroGPU':
-            generator = zerogpu_model_fn(prefix, history, gpu_pipe)
+            generator = zerogpu_model_fn(prefix, history, server)
         elif server == 'CPU':
-            generator = cpu_model_fn(prefix, history, cpu_pipe)
+            generator = cpu_model_fn(prefix, history, server)
         # elif server == 'RunPod':
         #     generator = runpod_model_fn(prefix, history)
         else:
@@ -186,8 +186,13 @@ CPUs will be slower but there is no time limit.
         for history in generator:
             yield history
 
-    def cpu_model_fn(prefix, history, pipe):
+    def cpu_model_fn(prefix, history, server):
         queue = Queue(maxsize=10)
+        if server == 'CPU':
+            pipe = cpu_pipe
+        else:
+            pipe = gpu_pipe
+            
         class MyStreamer:
             def put(self, tokens):
                 for token in tokens.flatten():