From 9bb9266845ac99baec00cbbd8a5da6a1c5a85e19 Mon Sep 17 00:00:00 2001 From: Dobidop <67412288+Dobidop@users.noreply.github.com> Date: Fri, 6 Feb 2026 11:16:33 +0100 Subject: [PATCH] Refactor imports and improve error handling in VR viewer components; update requirements for clarity --- GenerateStereo.py | 287 +----------------------------------------- __init__.py | 2 +- native_nodes.py | 7 +- native_viewer/core.py | 2 +- requirements.txt | 12 +- 5 files changed, 11 insertions(+), 299 deletions(-) diff --git a/GenerateStereo.py b/GenerateStereo.py index f250750..3d5be64 100644 --- a/GenerateStereo.py +++ b/GenerateStereo.py @@ -1,7 +1,5 @@ import torch import numpy as np -import os -import sys from typing import Union, List from PIL import Image import cv2 @@ -24,288 +22,10 @@ def log_memory(label=""): else: print(f"[MEM] {label}: RAM={ram_mb:.0f}MB") -file_dir = os.path.dirname(__file__) -sys.path.append(file_dir) - -import stereoimage_generation as sig +from . import stereoimage_generation as sig from comfy.utils import ProgressBar -import socket -import hashlib -import shutil -import time -import json -import psutil -import subprocess - - - - -class DeoVRLauncher: - def __init__(self, config_path): - self.config_path = config_path - self.deovr_path = self.load_deovr_path() - - def load_deovr_path(self): - """Loads the DeoVR installation path from a config file.""" - if not os.path.exists(self.config_path): - print(f"X Config file not found: {self.config_path}") - return None - - try: - with open(self.config_path, "r") as f: - config = json.load(f) - return config.get("deovr_path", None) - except Exception as e: - print(f"X Error reading config file: {e}") - return None - - def is_process_running(self, process_name): - """Checks if a process with the given name is running.""" - for process in psutil.process_iter(attrs=['name']): - if process.info['name'].lower() == process_name.lower(): - return True - return False - - def is_deovr_running(self): - """Checks if DeoVR.exe is currently running.""" - return self.is_process_running("deovr.exe") - - def create_dummy_image(self, temp_folder): - """Creates a black dummy image in the temp folder.""" - dummy_path = os.path.join(temp_folder, "dummy.png") - - if not os.path.exists(temp_folder): - os.makedirs(temp_folder) - - try: - img = Image.new("RGB", (1920, 1080), (0, 0, 0)) # Black image - img.save(dummy_path) - print(f"OK! Created dummy image: {dummy_path}") - return dummy_path - except Exception as e: - print(f"X Error creating dummy image: {e}") - return None - - def launch_deovr(self): - """Launches DeoVR and waits for it to start, checking every 2 seconds for up to 30 seconds.""" - - if self.is_deovr_running(): - print("OK! DeoVR is already running. No need to launch.") - return - - if not self.deovr_path: - print("XX DeoVR path is not set. Check the config file.") - return - - if not os.path.exists(self.deovr_path): - print("\nXXX DeoVR path seems to be incorrect, file not found:") - print(f"XXX {self.deovr_path}") - print("XXX Update the config file to point to your DeoVR.exe file to autostart it.") - print("!!! You can still use this manually by launching DeoVR and opening an image/video.") - return - - # Check if VR-related processes are running - vr_processes = ["vrmonitor.exe", "vrwebhelper.exe"] - initial_wait_time = 0 - - for proc in vr_processes: - if not self.is_process_running(proc): - print(f"! {proc} is not running. Waiting an extra 3 seconds.") - initial_wait_time += 3 - - # Try launching DeoVR - script_directory = os.path.dirname(os.path.abspath(__file__)) - loading_image_path = os.path.join(script_directory, 'loading_sbs_flat.png') - - try: - subprocess.Popen([self.deovr_path, loading_image_path], shell=True) - print(">> Launched DeoVR...") - except Exception as e: - print(f"XXX Error launching DeoVR: {e}") - return - - # Wait for DeoVR to start (polling every 2 seconds, max 30 seconds + extra time) - max_wait_time = 30 + initial_wait_time - elapsed_time = 1 - - print(" Waiting for DeoVR to start...") - - time.sleep(initial_wait_time) - - while elapsed_time < max_wait_time: - if self.is_deovr_running(): - print("OK! DeoVR has successfully launched.") - #time.sleep(10) # Final wait before proceeding - return - - if elapsed_time % 6 == 0: # Print status every 6 seconds - print(f" Still waiting... ({elapsed_time}/{max_wait_time} seconds)") - time.sleep(2) - elapsed_time += 2 - - print("XXX ERROR: DeoVR did not start within the expected time.") - - - -class DeoVRViewNode: - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "file_name": ("STRING", {"forceInput": True}), # User-provided file or video path - "projection_type": ([ - "Flat 2D plane", - "Equirectangular - 180° FOV", - "Equirectangular - 360° FOV", - "Fisheye projection - 180° FOV", - "Fisheye projection - 190° FOV", - "Fisheye projection - 200° FOV with MKX lens correction", - "Fisheye projection - 220° FOV with VRCA lens correction" - ],), - "eye_location": (["Side-by-Side", "Top-Bottom"],), - "file_location": (["Output folder", "Input folder", "Other (provide full path)"],), - } - } - - RETURN_TYPES = () - FUNCTION = "send_path_to_deovr" - CATEGORY = "Output" - OUTPUT_NODE = True - - def send_path_to_deovr(self, file_name, projection_type, eye_location, file_location): - """ - Copies the file to ComfyUI's temp folder with a modified name based on the projection type and eye location. - A unique MD5 hash is computed for the file to ensure uniqueness. If the file has been copied before, - the existing copy is reused. The command is then sent to DeoVR with the new file path. - """ - config_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "./config.json")) - launcher = DeoVRLauncher(config_path) - launcher.launch_deovr() - - if file_location == "Output folder": - file_name = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../output/", file_name)) - - if file_location == "Input folder": - file_name = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../input/", file_name)) - - - if not os.path.exists(file_name): - print(f"Error: The specified file does not exist: {file_name}") - return () - - # Define the temp folder relative to ComfyUI (two levels up) - temp_folder = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../temp")) - - if not os.path.exists(temp_folder): - os.makedirs(temp_folder) - - # Compute a unique hash for the file (MD5 is fast for most files) - file_hash = self.compute_file_hash(file_name) - - # Get the base name and extension of the file - base_name = os.path.basename(file_name) - base, ext = os.path.splitext(base_name) - - # Define suffix mappings (inspired by generateJson.py) - projection_suffix_mapping = { - "Flat 2D plane": "_screen", - "Equirectangular - 180° FOV": "_180", - "Equirectangular - 360° FOV": "_360", - "Fisheye projection - 180° FOV": "_fisheye", - "Fisheye projection - 190° FOV": "_fisheye190", - "Fisheye projection - 200° FOV with MKX lens correction": "_mkx200", - "Fisheye projection - 220° FOV with VRCA lens correction": "_vrca220" - } - stereo_suffix_mapping = { - "Side-by-Side": "_SBS", - "Top-Bottom": "_TB" - } - - proj_suffix = projection_suffix_mapping.get(projection_type, "") - stereo_suffix = stereo_suffix_mapping.get(eye_location, "") - - # Create the new file name: [hash]_[base][proj_suffix][stereo_suffix][ext] - new_file_name = f"{file_hash}_{base}{proj_suffix}{stereo_suffix}{ext}" - new_file_path = os.path.join(temp_folder, new_file_name) - - # If the file doesn't already exist in the temp folder, copy it. - if not os.path.exists(new_file_path): - try: - shutil.copy2(file_name, new_file_path) - print(f"Copied file to: {new_file_path}") - except Exception as e: - print(f"Error copying file: {e}") - return () - - # Create the JSON command to send to DeoVR. - command = { - "path": new_file_path, - "screenType": proj_suffix[1:] if proj_suffix.startswith("_") else "flat", - "stereoMode": stereo_suffix[1:].lower() if stereo_suffix.startswith("_") else "sbs", - "is3d": True - } - - - max_retries = 10 - tries = 0 - # Send the command - while tries <= max_retries: - message, errormessage = self.send_command_to_deovr(command) - - if message == 'error': - #print(f"Retrying to send command to DeoVR") - time.sleep(2) - tries = tries + 1 - else: - return () - - print(f"Failed to send command to DeoVR") - print(f"Error message: {errormessage}") - return () - - def send_command_to_deovr(self, command): - """Handles sending a command to the DeoVR remote control API and returns the response.""" - try: - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - s.connect(("127.0.0.1", 23554)) # DeoVR's remote control port - command_json = json.dumps(command) - length_prefix = len(command_json).to_bytes(4, byteorder='little') - packet = length_prefix + command_json.encode('utf-8') - - # Send command - s.sendall(packet) - #print(f"✅ Sent command to DeoVR: {command}") - - # Read response from server - response_length_bytes = s.recv(4) # First 4 bytes indicate message length - if not response_length_bytes: - #print("❌ No response received from DeoVR.") - return None - - response_length = int.from_bytes(response_length_bytes, byteorder='little') - response_data = s.recv(response_length).decode('utf-8') - response_json = json.loads(response_data) - - #print(f"📩 Received response from DeoVR: {response_json}") - return response_json, '' # Return the parsed response - - except Exception as e: - #print(f"❌ Error sending command to DeoVR: {e}") - return 'error', e - - def compute_file_hash(self, file_path): - """Computes an MD5 hash of the file for a unique identifier.""" - hash_md5 = hashlib.md5() - try: - with open(file_path, "rb") as f: - for chunk in iter(lambda: f.read(4096), b""): - hash_md5.update(chunk) - return hash_md5.hexdigest() - except Exception as e: - print(f"Error computing file hash: {e}") - return "errorhash" @@ -470,7 +190,8 @@ class StereoImageNode: depthmap_right_batch.append(np2tensor(right_depth_np)) # Generate mask (for GPU warp, there are no unfilled pixels due to border padding) - mask = torch.zeros((1, results_tensors[0].shape[1], results_tensors[0].shape[2] // 2)) + # Use original single-eye dimensions (correct for all modes: SBS, TB, anaglyph) + mask = torch.zeros((1, img_tensor.shape[1], img_tensor.shape[2])) mask_batch.append(mask) # Memory management for large batches - concatenate and clear @@ -755,10 +476,8 @@ def tensor2pil(image): NODE_CLASS_MAPPINGS = { "StereoImageNode": StereoImageNode, - "DeoVRViewNode": DeoVRViewNode } NODE_DISPLAY_NAME_MAPPINGS = { "StereoImageNode": "Stereo Image Node", - "DeoVRViewNode": "DeoVR View" } diff --git a/__init__.py b/__init__.py index a47fea4..700c350 100644 --- a/__init__.py +++ b/__init__.py @@ -23,7 +23,7 @@ except ImportError as e: print("\n" + "="*60) print("PyOpenXR not available. Native VR viewer nodes disabled.") print("To enable native VR viewing:") - print(" pip install -r requirements-native.txt") + print(" pip install -r requirements.txt") print("="*60 + "\n") # Import StereoDiffusion nodes (optional, requires diffusers) diff --git a/native_nodes.py b/native_nodes.py index 6fef61d..b1baf47 100644 --- a/native_nodes.py +++ b/native_nodes.py @@ -3,14 +3,12 @@ Native VR viewer nodes for ComfyUI using PyOpenXR Auto-launches directly into VR headset without browser """ +import hashlib import torch import numpy as np from PIL import Image -import io import os -import json import folder_paths -import sys from .native_viewer import ( check_openxr_available, @@ -106,7 +104,8 @@ class NativeStereoImageViewer: stereo_format = "Side-by-Side" # Save the stereo image - filename = f"{self.prefix_append}_{hash(str(img_np.tobytes()))}.png" + img_hash = hashlib.md5(img_np.tobytes()).hexdigest()[:16] + filename = f"{self.prefix_append}_{img_hash}.png" filepath = os.path.join(self.output_dir, filename) pil_img = Image.fromarray(img_np) diff --git a/native_viewer/core.py b/native_viewer/core.py index 7a4d4e4..1a962c7 100644 --- a/native_viewer/core.py +++ b/native_viewer/core.py @@ -636,7 +636,7 @@ class PersistentNativeViewer: while not self.media_queue.empty(): try: self.media_queue.get_nowait() - except: + except Exception: break # Mark as not running diff --git a/requirements.txt b/requirements.txt index 3164b3c..6a9c648 100644 --- a/requirements.txt +++ b/requirements.txt @@ -7,13 +7,8 @@ numba scipy psutil -# Note: For additional features, install optional dependencies: -# - Native VR Viewer: pip install -r requirements-native.txt -# - StereoDiffusion: pip install -r requirements-diffusion.txt - # Native VR Viewer Dependencies (PyOpenXR) -# Install these for native VR viewing with auto-launch to headset -# Usage: pip install -r requirements-native.txt +# For native VR viewing with auto-launch to headset pyopenxr>=1.0.0 PyOpenGL>=3.1.0 @@ -28,8 +23,7 @@ pygame>=2.0.0 # For audio playback # Mac: brew install ffmpeg # StereoDiffusion Dependencies -# Install these for AI-powered stereo generation using diffusion models -# Usage: pip install -r requirements-diffusion.txt +# For AI-powered stereo generation using diffusion models diffusers>=0.21.0 transformers @@ -38,4 +32,4 @@ einops tqdm scikit-image -# Note: torch, numpy, Pillow, and opencv-python are already included in base requirements.txt +# Note: torch, numpy, Pillow, and opencv-python are already included in base dependencies above