update gradio_web_server.py and model_worker.py

gradio_web_server.py

@@ -818,7 +818,7 @@ if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument("--host", type=str, default="0.0.0.0")
     parser.add_argument("--port", type=int, default=11000)
-    parser.add_argument("--controller-url", type=str, default="http://localhost:21001")
+    parser.add_argument("--controller-url", type=str, default=None)
     parser.add_argument("--concurrency-count", type=int, default=10)
     parser.add_argument(
         "--model-list-mode", type=str, default="once", choices=["once", "reload"]
@@ -829,6 +829,11 @@ if __name__ == "__main__":
     parser.add_argument("--embed", action="store_true")
     args = parser.parse_args()
     logger.info(f"args: {args}")
+    if not args.controller_url:
+        args.controller_url = os.environ.get("CONTROLLER_URL", None)
+
+    if not args.controller_url:
+        raise ValueError("controller-url is required.")
 
     models = get_model_list()
 

model_worker.py

@@ -160,6 +160,25 @@ def split_model(model_name):
     return device_map
 
 
+@spaces.GPU(duration=120)
+def multi_thread_infer(
+    model, tokenizer, pixel_values, question, history, generation_config
+):
+    with torch.no_grad():
+        thread = Thread(
+            target=model.chat,
+            kwargs=dict(
+                tokenizer=tokenizer,
+                pixel_values=pixel_values,
+                question=question,
+                history=history,
+                return_history=False,
+                generation_config=generation_config,
+            ),
+        )
+        thread.start()
+
+
 class ModelWorker:
     def __init__(
         self,
@@ -325,8 +344,6 @@ class ModelWorker:
             "queue_length": self.get_queue_length(),
         }
 
-    # @torch.inference_mode()
-    @spaces.GPU(duration=120)
     def generate_stream(self, params):
         system_message = params["prompt"][0]["content"]
         send_messages = params["prompt"][1:]
@@ -428,20 +445,14 @@ class ModelWorker:
             streamer=streamer,
         )
         logger.info(f"Generation config: {generation_config}")
-
-        with torch.no_grad():
-            thread = Thread(
-                target=self.model.chat,
-                kwargs=dict(
-                    tokenizer=self.tokenizer,
-                    pixel_values=pixel_values,
-                    question=question,
-                    history=history,
-                    return_history=False,
-                    generation_config=generation_config,
-                ),
-            )
-            thread.start()
+        multi_thread_infer(
+            self.model,
+            self.tokenizer,
+            pixel_values,
+            question,
+            history,
+            generation_config,
+        )
 
         generated_text = ""
         for new_text in streamer:
@@ -541,4 +552,4 @@ if __name__ == "__main__":
         args.load_8bit,
         args.device,
     )
-    uvicorn.run(app, host=args.host, port=args.port, log_level="info", workers=1)
+    uvicorn.run(app, host=args.host, port=args.port, log_level="info")