diff --git a/app.py b/app.py deleted file mode 100644 index 2ae5850..0000000 --- a/app.py +++ /dev/null @@ -1,200 +0,0 @@ -import os -import sys -import time - -import torch -from diffusers import AutoencoderTiny, StableDiffusionPipeline -from diffusers.utils import load_image - -sys.path.insert(0, os.path.abspath('../StreamDiffusion')) - -from streamdiffusion import StreamDiffusion -from streamdiffusion.image_utils import postprocess_image - -from utils.viewer import receive_images - -from utils.wrapper import StreamDiffusionWrapper -from threading import Thread - - -from multiprocessing import Process, Queue, get_context - -from perlin import perlin_2d, rand_perlin_2d, rand_perlin_2d_octaves, perlin_2d_octaves -from scene_prompt import surreal_prompt_parts -from scene_prompt import surreal_prompts -from scene_prompt import regret_prompts - -from spout_util import send_spout_image, get_spout_image - -from osc import start_osc_server - - -import fire - -def image_generation_process( - queue: Queue, - fps_queue: Queue, - prompt_queue: Queue, - input_queue: Queue, - # prompt: str, - model_id_or_path: str, -)-> None: - # stream = StreamDiffusionWrapper( - # model_id_or_path=model_id_or_path, - # lora_dict=None, - # t_index_list=[0, 16, 32, 45], - # frame_buffer_size=1, - # width=512, - # height=512, - # warmup=10, - # acceleration="xformers", - # mode="txt2img", - # use_denoising_batch=False, - # cfg_type="none", - # seed=2, - # ) - stream = StreamDiffusionWrapper( - model_id_or_path=model_id_or_path, - t_index_list=[0], - frame_buffer_size=1, - warmup=10, - acceleration="tensorrt", - use_lcm_lora=False, - mode="img2img", - cfg_type="none", - use_denoising_batch=True, - output_type="pil", - ) - - start_prompt = "A glowing, vintage phone booth standing in surreal landscapes across different scene" - # Prepare the stream - stream.prepare( - prompt=start_prompt, - num_inference_steps=50, - ) - - # Prepare image - # init_image = load_image("example.png").resize((512, 512)) - - # Warmup >= len(t_index_list) x frame_buffer_size - # for _ in range(stream.batch_size - 1): - # stream() - - previous_output = None - idx=0 - last_time = time.time() - - while True: - # try: - start_time = time.time() - # x_output = stream(image=previous_output) - # x_output=stream.stream.txt2img_sd_turbo(1).cpu() - - - - input_image= input_queue.get(block=True) - # input_image = stream.preprocess_image('input.png') - - - # Check if a new prompt is available in the prompt_queue - if not prompt_queue.empty(): - new_prompt = prompt_queue.get(block=False) - if new_prompt: - x_output = stream.img2img(image=input_image, prompt=new_prompt) - print(f"Received new prompt from queue: {new_prompt}") - else: - # Use the current prompt if no new prompt is available - x_output = stream.img2img(image=input_image) - - - - # preprocessed_image =stream.postprocess_image(x_output) - - queue.put(x_output, block=False) - - # queue.put(preprocessed_image, block=False) - - # Calculate FPS - elapsed_time = time.time() - start_time - fps = 1 / elapsed_time if elapsed_time > 0 else float('inf') - fps_queue.put(fps) - - # x_output = (x_output + 1) / 2 # Scale from [-1, 1] to [0, 1] - # x_output = torch.clamp(x_output, 0, 1) - # previous_output = x_output - - # except KeyboardInterrupt: - # print(f"fps: {fps}") - # return - - -def main()-> None: - - try: - ctx = get_context('spawn') - queue = Queue() - fps_queue = Queue() - # noise_queue = Queue() - spout_in_queue = Queue() - - # prompt = "A surreal landscapes" - # prompt=regret_prompts[0] - - prompt_queue = Queue() - - # model_id_or_path = "KBlueLeaf/kohaku-v2.1" - model_id_or_path = "stabilityai/sd-turbo" - - - # start_osc_server(prompt_queue) - process_osc = ctx.Process( - target=start_osc_server, - args=(prompt_queue,) - ) - process_osc.start() - - print("Starting spout input process") - process_spout_in = ctx.Process( - target=get_spout_image, - args=(spout_in_queue, 512, 512), - ) - process_spout_in.start() - - - print("Starting image generation process") - process_gen= ctx.Process( - target=image_generation_process, - args=(queue, fps_queue, prompt_queue, spout_in_queue, model_id_or_path), - ) - process_gen.start() - - - - # process_show=ctx.Process(target=receive_images, args=(queue, fps_queue)) - # process_show.start() - - # print("Starting spout output process") - process_spout_out=ctx.Process(target=send_spout_image, args=(queue, 512, 512)) - process_spout_out.start() - - - process_gen.join() - # process_spout_in.join() - process_spout_out.join() - process_osc.join() - - - except KeyboardInterrupt: - print("Process interrupted") - - process_gen.terminate() - # process_spout_in.terminate() - process_spout_out.terminate() - process_osc.terminate() - - return - - - -if __name__ == "__main__": - fire.Fire(main) diff --git a/input.png b/input.png deleted file mode 100644 index 3d6f015..0000000 Binary files a/input.png and /dev/null differ diff --git a/osc.py b/osc.py deleted file mode 100644 index 96b1dd8..0000000 --- a/osc.py +++ /dev/null @@ -1,23 +0,0 @@ -import argparse -import math - -from pythonosc.dispatcher import Dispatcher -from pythonosc import osc_server - -OSC_PORT = 8787 - - -def start_osc_server(queue): - - def onReceivePrompt(address, *args): - prompt = " ".join(args) - print(f"Received prompt: {prompt}") - queue.put(prompt) - - dispatcher = Dispatcher() - dispatcher.map("/prompt", onReceivePrompt) - - - server = osc_server.ThreadingOSCUDPServer(("localhost", OSC_PORT), dispatcher) - print(f"OSC server is running on port {OSC_PORT}") - server.serve_forever() \ No newline at end of file diff --git a/perlin.py b/perlin.py deleted file mode 100644 index be51937..0000000 --- a/perlin.py +++ /dev/null @@ -1,69 +0,0 @@ -import torch -import math - -def rand_perlin_2d(shape, res, fade = lambda t: 6*t**5 - 15*t**4 + 10*t**3): - delta = (res[0] / shape[0], res[1] / shape[1]) - d = (shape[0] // res[0], shape[1] // res[1]) - - grid = torch.stack(torch.meshgrid(torch.arange(0, res[0], delta[0]), torch.arange(0, res[1], delta[1])), dim = -1) % 1 - angles = 2*math.pi*torch.rand(res[0]+1, res[1]+1) - gradients = torch.stack((torch.cos(angles), torch.sin(angles)), dim = -1) - - tile_grads = lambda slice1, slice2: gradients[slice1[0]:slice1[1], slice2[0]:slice2[1]].repeat_interleave(d[0], 0).repeat_interleave(d[1], 1) - dot = lambda grad, shift: (torch.stack((grid[:shape[0],:shape[1],0] + shift[0], grid[:shape[0],:shape[1], 1] + shift[1] ), dim = -1) * grad[:shape[0], :shape[1]]).sum(dim = -1) - - n00 = dot(tile_grads([0, -1], [0, -1]), [0, 0]) - n10 = dot(tile_grads([1, None], [0, -1]), [-1, 0]) - n01 = dot(tile_grads([0, -1],[1, None]), [0, -1]) - n11 = dot(tile_grads([1, None], [1, None]), [-1,-1]) - t = fade(grid[:shape[0], :shape[1]]) - return math.sqrt(2) * torch.lerp(torch.lerp(n00, n10, t[..., 0]), torch.lerp(n01, n11, t[..., 0]), t[..., 1]) - -def rand_perlin_2d_octaves(shape, res, octaves=1, persistence=0.5): - noise = torch.zeros(shape) - frequency = 1 - amplitude = 1 - for _ in range(octaves): - noise += amplitude * rand_perlin_2d(shape, (frequency*res[0], frequency*res[1])) - frequency *= 2 - amplitude *= persistence - return noise - -def perlin_2d(shape, res, seed, fade=lambda t: 6*t**5 - 15*t**4 + 10*t**3): - delta = (res[0] / shape[0], res[1] / shape[1]) - d = (shape[0] // res[0], shape[1] // res[1]) - - grid = torch.stack(torch.meshgrid(torch.arange(0, res[0], delta[0]), torch.arange(0, res[1], delta[1])), dim=-1) % 1 - base_seed = int(seed) - frac_seed = seed - base_seed - - torch.manual_seed(base_seed) - angles_base = 2 * math.pi * torch.rand(res[0] + 1, res[1] + 1) - gradients_base = torch.stack((torch.cos(angles_base), torch.sin(angles_base)), dim=-1) - - torch.manual_seed(base_seed + 1) - angles_next = 2 * math.pi * torch.rand(res[0] + 1, res[1] + 1) - gradients_next = torch.stack((torch.cos(angles_next), torch.sin(angles_next)), dim=-1) - - gradients = (1 - frac_seed) * gradients_base + frac_seed * gradients_next - - tile_grads = lambda slice1, slice2: gradients[slice1[0]:slice1[1], slice2[0]:slice2[1]].repeat_interleave(d[0], 0).repeat_interleave(d[1], 1) - dot = lambda grad, shift: (torch.stack((grid[:shape[0], :shape[1], 0] + shift[0], grid[:shape[0], :shape[1], 1] + shift[1]), dim=-1) * grad[:shape[0], :shape[1]]).sum(dim=-1) - - n00 = dot(tile_grads([0, -1], [0, -1]), [0, 0]) - n10 = dot(tile_grads([1, None], [0, -1]), [-1, 0]) - n01 = dot(tile_grads([0, -1], [1, None]), [0, -1]) - n11 = dot(tile_grads([1, None], [1, None]), [-1, -1]) - t = fade(grid[:shape[0], :shape[1]]) - return math.sqrt(2) * torch.lerp(torch.lerp(n00, n10, t[..., 0]), torch.lerp(n01, n11, t[..., 0]), t[..., 1]) - - -def perlin_2d_octaves(shape, res, seed, octaves=1, persistence=0.5, fade=lambda t: 6*t**5 - 15*t**4 + 10*t**3): - noise = torch.zeros(shape) - frequency = 1 - amplitude = 1 - for i in range(octaves): - noise += amplitude * perlin_2d(shape, (frequency * res[0], frequency * res[1]), seed + i, fade) - frequency *= 2 - amplitude *= persistence - return noise \ No newline at end of file diff --git a/scene_prompt.py b/scene_prompt.py deleted file mode 100644 index e88ce2d..0000000 --- a/scene_prompt.py +++ /dev/null @@ -1,42 +0,0 @@ -surreal_prompts = [ - "a surreal landscape of floating islands under a glowing sky", - "an ethereal valley where waterfalls rise into the clouds", - "a dreamlike desert with mirrored sand and hovering stones", - "an endless ocean reflecting fractured moons and stars", - "a neon-lit canyon with levitating ruins and glowing mist", - "a twilight forest where the trees grow upside-down", - "a luminous terrain with bioluminescent plants and crystal arches", - "a gravity-defying mountain range spiraling into the void", - "a shattered realm of glass bridges and hovering towers", - "an alien world lit by pulsating constellations and fluid geometry" -] - -surreal_prompt_parts = [ - "a surreal landscape", - "with floating islands", - "glowing waterfalls", - "neon-colored skies", - "mirror-like desert ground", - "levitating rocks", - "upside-down trees", - "ancient ruins suspended in air", - "bioluminescent flora", - "shattered moons overhead", - "a path of glass tiles", - "crystal towers emitting soft hums", - "gravity-defying rivers", - "alien constellations glowing brightly" -] - -regret_prompts = [ - "a lone figure standing in a vast, empty desert at dusk", - "fractured mirrors scattered across the sand, reflecting different memories", - "a withered tree growing upside down from the sky, its roots dripping ink", - "floating clocks melting into the horizon, ticking backwards", - "ghostly silhouettes walking in reverse, retracing forgotten steps", - "a house half-submerged in water, its windows glowing faintly with past laughter", - "the sky opens into a tunnel of old photographs slowly burning at the edges", - "giant stone hands reaching out from the earth, trying to grasp something lost", - "an ocean made of letters never sent, waves crashing with whispered apologies", - "a child version of the figure stands alone, staring at the adult with distant eyes" -] diff --git a/spout_util.py b/spout_util.py deleted file mode 100644 index a0aa628..0000000 --- a/spout_util.py +++ /dev/null @@ -1,125 +0,0 @@ -import torch -import SpoutGL - -from itertools import islice, cycle, repeat -import array -from random import randint -import time -from OpenGL import GL - -from multiprocessing import Queue -import numpy as np -from PIL import Image - - -TARGET_FPS = 30 -SEND_WIDTH = 512 -SEND_HEIGHT = 512 - -alpha_cache = np.full((512, 512, 1), 255, dtype=np.uint8) - - - -def spout_buffer_to_tensor(buffer, width, height): - # np_buffer = np.asarray(buffer, dtype=np.uint8) - np_buffer=np.frombuffer(buffer, dtype=np.uint8) - image_bgra = np_buffer.reshape((height, width, 4)) - - image_rgb = image_bgra[..., [2, 1, 0]] - image_float = image_rgb.astype(np.float32) / 255.0 - # image_normalized = (image_float * 2.0) - 1.0 - tensor = torch.from_numpy(image_float).permute(2, 0, 1) - - del np_buffer # Free memory - del image_bgra # Free memory - del image_rgb # Free memory - del image_float # Free memory - - - return tensor.unsqueeze(0) - - -def get_spout_image(queue, wwidth: int, wheight: int) -> None: - with SpoutGL.SpoutReceiver() as receiver: - receiver.setReceiverName("Spout DX11 Sender") - image_bgra = np.zeros((SEND_HEIGHT, SEND_WIDTH, 4), dtype=np.uint8) - - while True: - result = receiver.receiveImage(image_bgra, GL.GL_RGBA, False, 0) - # print("Receive result", result) - - if receiver.isUpdated(): - continue - # width = receiver.getSenderWidth() - # height = receiver.getSenderHeight() - # image_bgra = array.array('B', [0] * (width * height * 4)) # Correctly reallocate buffer with updated size - # print("Spout Receiver updated, Buffer size", width, height) - - # if buffer and result and not SpoutGL.helpers.isBufferEmpty(buffer): - if SpoutGL.helpers.isBufferEmpty(image_bgra): - continue - # pixels=spout_buffer_to_tensor(buffer, width, height) - # print("get_spout_image", pixels.shape) - - image_rgb_array= image_bgra[:, :, [2, 1, 0]] - pixels=Image.fromarray(image_rgb_array, 'RGB') - queue.put(pixels, block=False) - - - # Wait until the next frame is ready - # Wait time is in milliseconds; note that 0 will return immediately - # receiver.waitFrameSync("SpoutSender", 10000) - - - - -def randcolor(): - return randint(0, 255) - - -def tensor_to_spout_image(tensor): - image = tensor.squeeze(0) - if image.device.type != "cpu": - image = image.cpu() - image = image.permute(1, 2, 0).numpy() - - if image.min() < 0: - image = (image + 1) / 2 # Scale from [-1, 1] to [0, 1] - image = np.clip(image * 255, 0, 255).astype(np.uint8) - - # h, w, _ = image_np.shape - # alpha = np.full((h, w, 1), 255, dtype=np.uint8) - image_rgba = np.concatenate((image, alpha_cache), axis=-1) - image_bgra = image_rgba[..., [2, 1, 0, 3]] - - del image # Free memory - - return np.ascontiguousarray(image_bgra) # Ensure the array is contiguous in memory - -def send_spout_image(queue: Queue, width: int, height: int)->None: - - with SpoutGL.SpoutSender() as sender: - sender.setSenderName("StreamDiffusion") - - while True: - - # Check if there are images in the queue - if not queue.empty(): - output_image = queue.get(block=False) - # pixels = tensor_to_spout_image(image) - - output_bgr_array = np.array(output_image, dtype=np.uint8)[:, :, ::-1] - output_bgra_array = np.zeros((SEND_HEIGHT, SEND_WIDTH, 4), dtype=np.uint8) - output_bgra_array[:, :, :3] = output_bgr_array - output_bgra_array[:, :, 3] = 255 - buffer = output_bgra_array - - - result = sender.sendImage(buffer, width, height, GL.GL_RGBA, False, 0) - # print("Send result", result) - - # Indicate that a frame is ready to read - sender.setFrameSync("StreamDiffusion") - - # Wait for next send attempt - # time.sleep(1./TARGET_FPS) \ No newline at end of file diff --git a/utils/viewer.py b/utils/viewer.py deleted file mode 100644 index dd6f6ca..0000000 --- a/utils/viewer.py +++ /dev/null @@ -1,98 +0,0 @@ -import os -import sys -import threading -import time -import tkinter as tk -from multiprocessing import Queue -from typing import List -from PIL import Image, ImageTk -from streamdiffusion.image_utils import postprocess_image - -sys.path.append(os.path.join(os.path.dirname(__file__), "..", "..")) - - -def update_image(image_data: Image.Image, label: tk.Label) -> None: - """ - Update the image displayed on a Tkinter label. - - Parameters - ---------- - image_data : Image.Image - The image to be displayed. - label : tk.Label - The labels where the image will be updated. - """ - width = 512 - height = 512 - tk_image = ImageTk.PhotoImage(image_data, size=width) - label.configure(image=tk_image, width=width, height=height) - label.image = tk_image # keep a reference - -def _receive_images( - queue: Queue, fps_queue: Queue, label: tk.Label, fps_label: tk.Label -) -> None: - """ - Continuously receive images from a queue and update the labels. - - Parameters - ---------- - queue : Queue - The queue to receive images from. - fps_queue : Queue - The queue to put the calculated fps. - label : tk.Label - The label to update with images. - fps_label : tk.Label - The label to show fps. - """ - while True: - try: - if not queue.empty(): - label.after( - 0, - update_image, - postprocess_image(queue.get(block=False), output_type="pil")[0], - label, - ) - if not fps_queue.empty(): - fps_label.config(text=f"FPS: {fps_queue.get(block=False):.2f}") - - time.sleep(0.0005) - except KeyboardInterrupt: - return - - -def receive_images(queue: Queue, fps_queue: Queue) -> None: - """ - Setup the Tkinter window and start the thread to receive images. - - Parameters - ---------- - queue : Queue - The queue to receive images from. - fps_queue : Queue - The queue to put the calculated fps. - """ - root = tk.Tk() - root.title("Image Viewer") - label = tk.Label(root) - fps_label = tk.Label(root, text="FPS: 0") - label.grid(column=0) - fps_label.grid(column=1) - - def on_closing(): - print("window closed") - root.quit() # stop event loop - return - - thread = threading.Thread( - target=_receive_images, args=(queue, fps_queue, label, fps_label), daemon=True - ) - thread.start() - - try: - root.protocol("WM_DELETE_WINDOW", on_closing) - root.mainloop() - except KeyboardInterrupt: - return -