[examples] fix multiprocessing on Linux

- use multiprocessing context to specify the spawn start method, which fixes the "RuntimeError: Cannot re-initialize CUDA in forked subprocess" on Linux (verified with Ubuntu 22.04 and kernel 6.2.0-39)
- call `.join()` to wait for processes to complete (avoids exiting program immediately)

examples/optimal-performance/multi.py

@@ -3,7 +3,7 @@
 import threading
 import time
 import tkinter as tk
-from multiprocessing import Process, Queue
+from multiprocessing import Process, Queue, get_context
 from typing import List, Literal
 
 import fire
@@ -174,17 +174,20 @@ def main(
  """
  Main function to start the image generation and viewer processes.
  """
- queue = Queue()
- fps_queue = Queue()
- process1 = Process(
+ ctx = get_context('spawn')
+ queue = ctx.Queue()
+ fps_queue = ctx.Queue()
+ process1 = ctx.Process(
  target=image_generation_process,
  args=(queue, fps_queue, prompt, model_id_or_path, batch_size, acceleration),
  )
  process1.start()
 
- process2 = Process(target=receive_images, args=(queue, fps_queue))
+ process2 = ctx.Process(target=receive_images, args=(queue, fps_queue))
  process2.start()
 
+ process1.join()
+ process2.join()
 
 if __name__ == "__main__":
  fire.Fire(main)

examples/optimal-performance/single.py

@@ -1,7 +1,7 @@
 import os
 import sys
 import time
-from multiprocessing import Process, Queue
+from multiprocessing import Process, Queue, get_context
 from typing import Literal
 
 import fire
@@ -72,17 +72,20 @@ def main(
  """
  Main function to start the image generation and viewer processes.
  """
- queue = Queue()
- fps_queue = Queue()
- process1 = Process(
+ ctx = get_context('spawn')
+ queue = ctx.Queue()
+ fps_queue = ctx.Queue()
+ process1 = ctx.Process(
  target=image_generation_process,
  args=(queue, fps_queue, prompt, model_id_or_path, acceleration),
  )
  process1.start()
 
- process2 = Process(target=receive_images, args=(queue, fps_queue))
+ process2 = ctx.Process(target=receive_images, args=(queue, fps_queue))
  process2.start()
 
+ process1.join()
+ process2.join()
 
 if __name__ == "__main__":
  fire.Fire(main)

examples/screen/main.py

@@ -2,7 +2,7 @@
 import sys
 import time
 import threading
-from multiprocessing import Process, Queue
+from multiprocessing import Process, Queue, get_context
 from typing import List, Literal, Dict, Optional
 import torch
 import PIL.Image
@@ -216,10 +216,10 @@ def main(
  Main function to start the image generation and viewer processes.
  """
  monitor = dummy_screen(width, height)
-
- queue = Queue()
- fps_queue = Queue()
- process1 = Process(
+ ctx = get_context('spawn')
+ queue = ctx.Queue()
+ fps_queue = ctx.Queue()
+ process1 = ctx.Process(
  target=image_generation_process,
  args=(
  queue,
@@ -246,9 +246,11 @@ def main(
  )
  process1.start()
 
- process2 = Process(target=receive_images, args=(queue, fps_queue))
+ process2 = ctx.Process(target=receive_images, args=(queue, fps_queue))
  process2.start()
 
+ process1.join()
+ process2.join()
 
 if __name__ == "__main__":
  fire.Fire(main)