Better init event waiting

app.py

@@ -306,7 +306,6 @@ if __name__ == "__main__":
     args.vqgan_config_name = "vqgan_pretrain"
 
     logger.info("Loading Llama model...")
-    init_event = threading.Event()
     llama_queue = launch_thread_safe_queue(
         config_name=args.llama_config_name,
         checkpoint_path=args.llama_checkpoint_path,
@@ -314,10 +313,8 @@ if __name__ == "__main__":
         precision=args.precision,
         max_length=args.max_length,
         compile=args.compile,
-        init_event=init_event,
     )
     llama_tokenizer = AutoTokenizer.from_pretrained(args.tokenizer)
-    init_event.wait()
     logger.info("Llama model loaded, loading VQ-GAN model...")
 
     vqgan_model = load_vqgan_model(

tools/llama/generate.py

@@ -600,6 +600,7 @@ def generate_long(
             yield all_codes
 
 
+
 def launch_thread_safe_queue(
     config_name,
     checkpoint_path,
@@ -607,17 +608,15 @@ def launch_thread_safe_queue(
     precision,
     max_length,
     compile=False,
-    init_event=None,
 ):
     input_queue = queue.Queue()
+    init_event = threading.Event()
 
     def worker():
         model, decode_one_token = load_model(
             config_name, checkpoint_path, device, precision, max_length, compile=compile
         )
-
-        if init_event is not None:
-            init_event.set()
+        init_event.set()
 
         while True:
             item = input_queue.get()
@@ -641,6 +640,7 @@ def launch_thread_safe_queue(
             event.set()
 
     threading.Thread(target=worker, daemon=True).start()
+    init_event.wait()
 
     return input_queue