在先前的 StreamDiffusion 環境建置文章中,我們是直接設定一組提示詞來生成圖片,每一次的生成都需要停止程式、修改、再重啟,非常不方便。
理想的目標是能從外部介面來做到隨時修改提示詞,因此我們需要透過 OSC 傳輸技術來接收程式外部的訊息。再者,目前生成的影像是直接儲存成 PNG 圖片,本文將會應用 NDI 影像傳輸技術,來將生成後的影像傳輸到程式外部,我們就能用 TouchDesigner 或 Unity 等工具來接收,再來做後續的互動設計。
安裝套件
我們沿用 StreamDiffusion 環境建置建立的開發環境,請使用 VS Code 終端機,輸入以下指令來安裝 OpenCV、OSC 與 NDI 的 Python 套件。
1uv add opencv-python python-osc ndi-python==5.1.1.1
即時文字生圖片
本文使用先前開發的文字生圖片程式碼來改寫,透過 OSC 來即時接收新的提示詞並更新,用 NDI 技術來即時輸出影像。
建立提示詞更新與影像輸出功能
我們需要大改程式,有動到的地方已於以下程式碼中標示。請編輯 main.py,並複製貼上以下程式碼:
1from utils.wrapper import StreamDiffusionWrapper
2from threading import Thread
3import cv2
4import NDIlib as ndi
5import numpy as np
6import pythonosc
7import pythonosc.dispatcher
8import pythonosc.osc_server
9
10def main():
11 # 初始化模型
12 global prompt
13 prompt = "cat, detailed, fantasy, 8k"
14 stream = StreamDiffusionWrapper(
15 model_id_or_path="stabilityai/sd-turbo",
16 t_index_list=[0, 24, 32],
17 width=512,
18 height=512,
19 cfg_type="none",
20 acceleration="xformers",
21 mode="txt2img",
22 output_type="np" # 指定輸出影像為 numpy.ndarray
23 )
24 stream.prepare(prompt=prompt)
25
26 # 建立 OSC 連線
27 osc_dispatcher = pythonosc.dispatcher.Dispatcher()
28 osc_dispatcher.map("/prompt", lambda _, args: globals().update(prompt=args))
29 osc_server = pythonosc.osc_server.ThreadingOSCUDPServer(("127.0.0.1", 7000), osc_dispatcher)
30 osc_server_thread = Thread(target=osc_server.serve_forever)
31 osc_server_thread.start()
32
33 # 建立 NDI 連線
34 ndi.initialize()
35 ndi_send_cfg = ndi.SendCreate()
36 ndi_send_cfg.ndi_name = "ndi-output"
37 ndi_send = ndi.send_create(ndi_send_cfg)
38 ndi_send_frame = ndi.VideoFrameV2()
39 ndi_send_frame.FourCC = ndi.FOURCC_VIDEO_TYPE_RGBA
40
41 # 模型熱機
42 for _ in range(stream.batch_size - 1):
43 stream(prompt=prompt)
44
45 # 主循環
46 try:
47 while True:
48 # 生成影像
49 output = stream(prompt=prompt)
50
51 # 格式轉換:float32 -> int8
52 output = (output * 255).astype(np.uint8)
53
54 # 格式轉換:RGB -> RGBA
55 output = cv2.cvtColor(output, cv2.COLOR_RGB2RGBA)
56
57 # NDI 發送
58 ndi_send_frame.data = output
59 ndi.send_send_video_v2(ndi_send, ndi_send_frame)
60
61 except KeyboardInterrupt:
62 print("中斷執行")
63
64 except Exception as e:
65 print(e)
66
67 finally:
68 # 釋放資源
69 ndi.send_destroy(ndi_send)
70 ndi.destroy()
71 osc_server.shutdown()
72 osc_server_thread.join()
73 print("已停止")
74
75
76if __name__ == "__main__":
77 main()
實際執行結果
我們用 TouchDesigner 來快速測試 OSC 和 NDI 的傳輸功能,可以從下圖中看到 StreamDiffusion 會一直生成貓咪的圖片,直到我們按下按鈕之後才變成馬,代表 StreamDiffusion 有成功接收到提示詞,並改變了影像生成的結果。
程式開發說明
在上面的程式中,我們主要開發了這些功能:
- 新建立了 prompt 變數,來儲存生成影像用的提示詞。
- 指定 StreamDiffusion 輸出的影像格式是 numpy.ndarray。
- 建立 OSC 伺服器,在接收到提示詞的時候會更新 prompt 變數。
- 建立 NDI 發送器,會用「ndi-output」這個名稱來發送影像。
- 用一個 while 無限迴圈作為持續生成影像的主循環,在主循環裡面我們會根據 prompt 儲存的提示詞來生成影像,並將生成後的影像格式轉換成 NDI 可用的格式,最後送出。
即時圖片生圖片
圖片生圖片的功能和文字生圖片的差別在有沒有接收外部的影像輸入,我們可以直接修改即時文字生圖片的程式,用 NDI 來接收外部輸入的影像,再來進行即時圖片生圖片的功能。
建立影像接收功能
本文直接修改即時文字生圖片的程式碼,請編輯 main.py,並複製貼上以下程式碼:
1from utils.wrapper import StreamDiffusionWrapper
2from threading import Thread
3from PIL import Image
4import cv2
5import NDIlib as ndi
6import numpy as np
7import pythonosc
8import pythonosc.dispatcher
9import pythonosc.osc_server
10
11def main():
12 # 初始化模型
13 global prompt
14 prompt = "cat, detailed, fantasy, 8k"
15 stream = StreamDiffusionWrapper(
16 model_id_or_path="stabilityai/sd-turbo",
17 t_index_list=[24, 32],
18 width=512,
19 height=512,
20 acceleration="xformers",
21 mode="img2img",
22 output_type="np" # 指定輸出影像為 numpy.ndarray
23 )
24 stream.prepare(prompt=prompt)
25
26 # 建立 OSC 連線
27 osc_dispatcher = pythonosc.dispatcher.Dispatcher()
28 osc_dispatcher.map("/prompt", lambda _, args: globals().update(prompt=args))
29 osc_server = pythonosc.osc_server.ThreadingOSCUDPServer(("127.0.0.1", 7000), osc_dispatcher)
30 osc_server_thread = Thread(target=osc_server.serve_forever)
31 osc_server_thread.start()
32
33 # 建立 NDI 連線
34 ndi.initialize()
35 ndi_recv_find = ndi.find_create_v2()
36 ndi_recv_sources = []
37 while not len(ndi_recv_sources) > 0:
38 ndi.find_wait_for_sources(ndi_recv_find, 1000)
39 ndi_recv_sources = ndi.find_get_current_sources(ndi_recv_find)
40
41 ndi_recv_source = None
42 for s in ndi_recv_sources:
43 if s.ndi_name.endswith(f"(TouchDesigner)"):
44 ndi_recv_source = s
45 break
46
47 ndi_recv_cfg = ndi.RecvCreateV3()
48 ndi_recv_cfg.color_format = ndi.RECV_COLOR_FORMAT_BGRX_BGRA
49 ndi_recv = ndi.recv_create_v3(ndi_recv_cfg)
50 ndi.recv_connect(ndi_recv, ndi_recv_source)
51 ndi.find_destroy(ndi_recv_find)
52 ndi_send_cfg = ndi.SendCreate()
53 ndi_send_cfg.ndi_name = "ndi-output"
54 ndi_send = ndi.send_create(ndi_send_cfg)
55 ndi_send_frame = ndi.VideoFrameV2()
56 ndi_send_frame.FourCC = ndi.FOURCC_VIDEO_TYPE_RGBA
57
58 # 模型熱機
59 for _ in range(stream.batch_size - 1):
60 stream(prompt=prompt)
61
62 # 主循環
63 try:
64 while True:
65 # 接收 NDI
66 t, v, _, _ = ndi.recv_capture_v2(ndi_recv, 1000)
67 if t != ndi.FRAME_TYPE_VIDEO:
68 continue
69 input = np.copy(v.data)
70 input = Image.fromarray(input)
71 input = stream.preprocess_image(input)
72 ndi.recv_free_video_v2(ndi_recv, v)
73
74 # 生成影像
75 output = stream(prompt=prompt, image=input)
76
77 # 格式轉換:float32 -> int8
78 output = (output * 255).astype(np.uint8)
79
80 # 格式轉換:RGB -> RGBA
81 output = cv2.cvtColor(output, cv2.COLOR_RGB2RGBA)
82
83 # NDI 發送
84 ndi_send_frame.data = output
85 ndi.send_send_video_v2(ndi_send, ndi_send_frame)
86
87 except KeyboardInterrupt:
88 print("中斷執行")
89
90 except Exception as e:
91 print(e)
92
93 finally:
94 # 釋放資源
95 ndi.recv_destroy(ndi_recv)
96 ndi.send_destroy(ndi_send)
97 ndi.destroy()
98 osc_server.shutdown()
99 osc_server_thread.join()
100 print("已停止")
101
102
103if __name__ == "__main__":
104 main()
實際執行結果
一樣用 TouchDesigner 來測試,把一個上了隨機波動的圖片來當作 StreamDiffusion 的影像輸入,我們就能獲得即時生成的貓咪影像。
程式開發說明
在上面的程式中,我們主要開發了這些功能:
- 建立 NDI 的影像接收功能,在啟動 StreamDiffusion 的時候會自動找到用「TouchDesigner」名稱來輸出影像的 NDI 介面,並在主循環中接收影像。
- 在主循環中用 NDI 接收到影像後,進行影像前處理後再連同提示詞來生成圖片。
結語
小結一下 StreamDiffusion 的開發進度,目前我們已經完成環境建置與提示詞、影像傳輸功能的開發。接下來,我們將著手把這個即時影像生成功能來和 TouchDesigner 與 Unity 進行整合,實際的讓這個功能成為互動設計的一環。