Merge pull request cumulo-autumn#39 from GradientSurfer/main

[examples] fix multiprocessing on Linux

Author: cumulo-autumn

examples/README.md

@@ -10,7 +10,7 @@ If you want to maximize performance, you need to install with following steps ex
 
 ## `screen/`
 
-Take a screen capture and process it. **This script only works on Windows.**
+Take a screen capture and process it.
 
 When you run the script, a translucent window appears. Position it at where you want to capture the screen and press the enter key to finalize the capture area.
 

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)