[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)

Author: GradientSurfer

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)