Update app.py

app.py

@@ -24,6 +24,7 @@ pipe = transformers.pipeline(
     torch_dtype="bfloat16",
     device="cuda",
 )
+cpu_pipe = pipe.to("cpu")
 print('Done')
 
 example_prefix = '''pitch duration wait velocity instrument
@@ -166,9 +167,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)
+            generator = zerogpu_model_fn(prefix, history, pipe)
         elif server == "CPU":
-            generator = cpu_model_fn(prefix, history)
+            generator = cpu_model_fn(prefix, history, cpu_pipe)
         # elif server == "RunPod":
         #     generator = runpod_model_fn(prefix, history)
         else:
@@ -176,7 +177,7 @@ CPUs will be slower but there is no time limit.
         for history in generator:
             yield history
 
-    def cpu_model_fn(prefix, history):
+    def cpu_model_fn(prefix, history, pipe):
         queue = Queue(maxsize=10)
         class MyStreamer:
             def put(self, tokens):
@@ -188,15 +189,12 @@ CPUs will be slower but there is no time limit.
             def end(self):
                 queue.put(None)
         def background_fn():
-            try:
-                result = pipe(
-                    prefix,
-                    streamer=MyStreamer(), 
-                    max_new_tokens=500,
-                    top_p=0.9, temperature=0.6,
-                )
-            except queue.Full:
-                print("Queue is full. Exiting.")
+            result = pipe(
+                prefix,
+                streamer=MyStreamer(), 
+                max_new_tokens=500,
+                top_p=0.9, temperature=0.6,
+            )
             print('Generated text:')
             print(result[0]['generated_text'])
             print()