From b1ebd4aa8fba1d6a1f37bc395a96f40e8a861a8a Mon Sep 17 00:00:00 2001 From: John | Elite Encoder Date: Wed, 2 Apr 2025 17:49:43 -0400 Subject: [PATCH] Revert "Super resolution (#16)" (#17) This reverts commit 59035c455a8e19b5a2b37da6b279b4607882d5d6. Co-authored-by: Rick Staa --- .gitignore | 9 - node_wrappers/super_resolution.py | 16 -- requirements.txt | 3 - src/super_resolution/README.md | 88 ------ src/super_resolution/__init__.py | 1 - src/super_resolution/models/__init__.py | 5 - src/super_resolution/models/sr_model.py | 251 ------------------ src/super_resolution/models/test_opencv_sr.py | 110 -------- src/super_resolution/nodes/__init__.py | 1 - src/super_resolution/nodes/sr_nodes.py | 95 ------- src/super_resolution/utils/__init__.py | 2 - src/super_resolution/utils/model_utils.py | 149 ----------- 12 files changed, 730 deletions(-) delete mode 100644 node_wrappers/super_resolution.py delete mode 100644 src/super_resolution/README.md delete mode 100644 src/super_resolution/__init__.py delete mode 100644 src/super_resolution/models/__init__.py delete mode 100644 src/super_resolution/models/sr_model.py delete mode 100644 src/super_resolution/models/test_opencv_sr.py delete mode 100644 src/super_resolution/nodes/__init__.py delete mode 100644 src/super_resolution/nodes/sr_nodes.py delete mode 100644 src/super_resolution/utils/__init__.py delete mode 100644 src/super_resolution/utils/model_utils.py diff --git a/.gitignore b/.gitignore index a8e8e64..27425f4 100644 --- a/.gitignore +++ b/.gitignore @@ -5,12 +5,3 @@ __pycache__ # IDEs .vscode *.code-workspace - -# Model weights -*.pb -*.pth -*.safetensors -*.pt -*.pth.tar -*.pth.tar.gz -*.pth.tar.xz diff --git a/node_wrappers/super_resolution.py b/node_wrappers/super_resolution.py deleted file mode 100644 index c19e924..0000000 --- a/node_wrappers/super_resolution.py +++ /dev/null @@ -1,16 +0,0 @@ -# Import node classes -from ..src.super_resolution.nodes.sr_nodes import SuperResolutionModelLoader, SuperResolutionUpscale - -# Define NODE_CLASS_MAPPINGS for ComfyUI -NODE_CLASS_MAPPINGS = { - "SuperResolutionModelLoader": SuperResolutionModelLoader, - "SuperResolutionUpscale": SuperResolutionUpscale, -} - -# Define NODE_DISPLAY_NAME_MAPPINGS for ComfyUI -NODE_DISPLAY_NAME_MAPPINGS = { - "SuperResolutionModelLoader": "SR Model Loader", - "SuperResolutionUpscale": "SR Upscale", -} - -__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index bfb1173..5c7905c 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,6 +1,3 @@ mediapipe torch diffusers -numpy -requests -opencv-contrib-python>=4.8.0 diff --git a/src/super_resolution/README.md b/src/super_resolution/README.md deleted file mode 100644 index c6ef872..0000000 --- a/src/super_resolution/README.md +++ /dev/null @@ -1,88 +0,0 @@ -# ComfyUI Super Resolution - -A collection of high-performance neural network-based Super Resolution models for ComfyUI. - -## Features - -- **Multiple SR Models**: Choose from FSRCNN, ESPCN, LapSRN, and EDSR with different quality/speed tradeoffs -- **Modular Design**: Load models once and reuse them across multiple upscaling operations -- **CUDA Acceleration**: Automatic GPU acceleration when available -- **Multiple Scale Factors**: Support for 2x, 3x, and 4x upscaling - -## Installation - -1. Clone this repository to your ComfyUI custom_nodes directory: - ``` - cd ComfyUI/custom_nodes - git clone https://github.com/yourusername/ComfyUI_SuperResolution - ``` - -2. Restart ComfyUI - the necessary models will be automatically downloaded on first use - -## Usage - -1. Add a **SR Model Loader** node to your workflow -2. Select your desired model type and scale factor -3. Connect the model output to a **SR Upscale** node -4. Connect an image to the upscale node -5. Run your workflow to get high-quality upscaled images! - -## Model Comparison - -| Model | Architecture | Features | Best For | Speed | Quality | -|-------|-------------|----------|----------|-------|---------| -| **FSRCNN-small** | Lightweight CNN | Fast, minimal memory use | Real-time processing, mobile | ★★★★★ | ★★ | -| **FSRCNN** | CNN with larger features | Good balance of speed/quality | General purpose | ★★★★ | ★★★ | -| **ESPCN** | Sub-pixel convolutions | Efficient upscaling at end | Text/line drawings | ★★★★ | ★★★ | -| **VDSR** | Very deep CNN | Better edge reconstruction | Detailed images with edges | ★★★ | ★★★★ | -| **LapSRN** | Laplacian pyramid | Progressive upscaling | Sharp edges, details | ★★★ | ★★★★ | -| **EDSR** | Deep residual network | Most parameters, best quality | Maximum detail | ★★ | ★★★★★ | - -## Technical Details - -Each model has unique architectural characteristics: - -1. **FSRCNN / FSRCNN-small**: Direct mapping from low to high resolution, lightweight with fewer parameters. - -2. **ESPCN**: Uses the efficient sub-pixel convolution technique that processes at low resolution and only expands dimensions at the final layer. - -3. **VDSR**: Very Deep Super Resolution network with 20 convolutional layers, uses global residual learning. - -4. **LapSRN**: Uses a Laplacian pyramid structure to progressively upscale images, which preserves edges better than the others. - -5. **EDSR**: Enhanced Deep Super-Resolution, a residual network with significantly more parameters, offering the highest quality but slower processing. - -## Performance Expectations (RTX 4090) - -| Model | 2x Upscaling | 4x Upscaling | Memory Usage | -|-------|--------------|--------------|--------------| -| FSRCNN-small | 200+ FPS | 150+ FPS | < 100MB | -| FSRCNN | 100+ FPS | 80+ FPS | < 150MB | -| ESPCN | 90+ FPS | 70+ FPS | < 150MB | -| LapSRN | 60+ FPS | 40+ FPS | < 200MB | -| EDSR | 40+ FPS | 25+ FPS | < 500MB | - -## When to Use Each Model - -- **FSRCNN-small**: When maximum speed is required -- **FSRCNN**: For a good balance of quality and performance -- **ESPCN**: For text and line art upscaling -- **LapSRN**: When you need better edge preservation than FSRCNN -- **EDSR**: When maximum quality is desired and performance is secondary - -## Comparison to Traditional Methods - -These neural network-based upscaling methods offer significantly better quality than traditional algorithms like Lanczos, Bicubic, or Bilinear interpolation, while often being just as fast thanks to CUDA acceleration. - -## Credits - -This nodepack implements models originally created by: - -- FSRCNN: Dong et al. (https://github.com/ryanontheinside/FSRCNN_Tensorflow) -- EDSR: Lim et al. (https://github.com/ryanontheinside/EDSR_Tensorflow) -- ESPCN: Shi et al. (https://github.com/ryanontheinside/TF-ESPCN) -- LapSRN: Lai et al. (https://github.com/ryanontheinside/TF-LapSRN) - -## License - -MIT \ No newline at end of file diff --git a/src/super_resolution/__init__.py b/src/super_resolution/__init__.py deleted file mode 100644 index 496a4cb..0000000 --- a/src/super_resolution/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# Dummy file ensuring this sub-package will be recognized. \ No newline at end of file diff --git a/src/super_resolution/models/__init__.py b/src/super_resolution/models/__init__.py deleted file mode 100644 index e5dcbe2..0000000 --- a/src/super_resolution/models/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -# Makes models directory a proper Python package - -from .sr_model import SuperResolutionModel - -__all__ = ["SuperResolutionModel"] \ No newline at end of file diff --git a/src/super_resolution/models/sr_model.py b/src/super_resolution/models/sr_model.py deleted file mode 100644 index c584eed..0000000 --- a/src/super_resolution/models/sr_model.py +++ /dev/null @@ -1,251 +0,0 @@ -import os -import cv2 -import torch -import numpy as np - -class SuperResolutionModel: - """Unified Super Resolution Model for all algorithm types""" - - # Configuration registry for all model types - MODEL_CONFIGS = { - "fsrcnn": { - "algorithm": "fsrcnn", - "repo_subdir": "FSRCNN_Tensorflow/models", - "filename_pattern": "FSRCNN_x{scale}.pb", - "supported_scales": [2, 3, 4], - }, - "fsrcnn-small": { - "algorithm": "fsrcnn", - "repo_subdir": "FSRCNN_Tensorflow/models", - "filename_pattern": "FSRCNN-small_x{scale}.pb", - "supported_scales": [2, 3, 4], - }, - "edsr": { - "algorithm": "edsr", - "repo_subdir": "EDSR_Tensorflow/models", - "filename_pattern": "EDSR_x{scale}.pb", - "supported_scales": [2, 3, 4], - }, - "espcn": { - "algorithm": "espcn", - "repo_subdir": "TF-ESPCN/export", - "filename_pattern": "ESPCN_x{scale}.pb", - "supported_scales": [2, 3, 4], - }, - "lapsrn": { - "algorithm": "lapsrn", - "repo_subdir": "TF-LapSRN/export", - "filename_pattern": "LapSRN_x{scale}.pb", - "supported_scales": [2, 3, 4], - }, - "vdsr": { - "algorithm": "vdsr", - "repo_subdir": "TF-VDSR/export", - "filename_pattern": "VDSR_x{scale}.pb", - "supported_scales": [2, 3, 4], - } - } - - @classmethod - def is_cuda_available(cls): - """Check if CUDA is available in OpenCV""" - # Check if CUDA module exists in OpenCV - if not hasattr(cv2, "cuda") or cv2.cuda.getCudaEnabledDeviceCount() == 0: - return False - - # Check if DNN module has CUDA support - try: - # Try to get available backends if the method exists - if hasattr(cv2.dnn, "getAvailableBackends"): - backends = cv2.dnn.getAvailableBackends() - if cv2.dnn.DNN_BACKEND_CUDA not in backends: - return False - - # Check for CUDA targets if method exists - if hasattr(cv2.dnn, "getAvailableTargets"): - targets = cv2.dnn.getAvailableTargets(cv2.dnn.DNN_BACKEND_CUDA) - if len(targets) == 0: - return False - except Exception: - return False - - return True - - @classmethod - def create(cls, model_type, scale_factor, **kwargs): - """Factory method to create a super resolution model by type""" - model_type = model_type.lower() - - if model_type not in cls.MODEL_CONFIGS: - raise ValueError(f"Unknown model type: {model_type}") - - return cls(model_type, scale_factor, **kwargs) - - def __init__(self, model_name, scale_factor, device="cuda" if torch.cuda.is_available() else "cpu"): - self.model_name = model_name.lower() - self.scale_factor = int(scale_factor) - self.device = device - self.model = None - self.loaded = False - self.using_cuda = False - - # Get configuration for this model - self.config = self.MODEL_CONFIGS.get(self.model_name) - if not self.config: - raise ValueError(f"No configuration found for model: {model_name}") - - # Validate scale factor - if self.scale_factor not in self.config["supported_scales"]: - raise ValueError(f"Scale factor {scale_factor} not supported for {model_name}. " - f"Supported scales: {self.config['supported_scales']}") - - def load_model(self, model_path=None): - """Load the super resolution model""" - if model_path is None: - # Try to find model in standard locations - model_dir = self.model_name.upper() - - # Search paths from highest to lowest priority - search_paths = [ - os.path.join(os.path.dirname(os.path.abspath(__file__)), "../../downloaded_models", model_dir), - os.path.join(os.path.dirname(os.path.abspath(__file__)), "../../models"), - os.path.expanduser("~/.cache/comfyui/custom_nodes/ComfyUI_SuperResolution/models"), - ] - - # Add repo-specific path if available - if self.config["repo_subdir"]: - search_paths.append( - os.path.join(os.path.dirname(os.path.abspath(__file__)), "../../../", self.config["repo_subdir"]) - ) - - for path in search_paths: - filename = self.config["filename_pattern"].format(scale=self.scale_factor) - candidate = os.path.join(path, filename) - if os.path.exists(candidate): - model_path = candidate - break - - if not model_path or not os.path.exists(model_path): - raise FileNotFoundError(f"Could not find {self.model_name} x{self.scale_factor} model file") - - # Load model using OpenCV DNN Super Resolution with the correct method - try: - # Try different methods to create the DNN SuperRes object for cross-platform compatibility - # This handles differences between Windows and Linux OpenCV builds - try: - # Method 1: Windows-style API (most common) - self.model = cv2.dnn_superres.DnnSuperResImpl_create() - except AttributeError: - try: - # Method 2: More generic constructor if available - self.model = cv2.dnn_superres.SuperResolution_create() - except AttributeError: - # Method 3: Direct class instantiation if available - self.model = cv2.dnn_superres.DnnSuperResImpl() - - self.model.readModel(model_path) - self.model.setModel(self.config["algorithm"], self.scale_factor) - - # Check CUDA availability during model loading - # Note: We don't raise errors here, as that will be handled during upscale if use_cuda=True - system_has_cuda = torch.cuda.is_available() - opencv_has_cuda = self.is_cuda_available() - - if system_has_cuda and opencv_has_cuda: - self.model.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA) - self.model.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA) - print(f"Model loaded with CUDA backend for {self.model_name}") - self.using_cuda = True - else: - self.model.setPreferableBackend(cv2.dnn.DNN_BACKEND_DEFAULT) - self.model.setPreferableTarget(cv2.dnn.DNN_TARGET_CPU) - print(f"Model loaded with CPU backend for {self.model_name}") - self.using_cuda = False - - self.loaded = True - except Exception as e: - raise RuntimeError(f"Error with OpenCV DNN Super Resolution: {e}. Check that opencv-contrib-python is installed.") - - return self - - def upscale(self, image, use_cuda=True): - """Upscale an image using the model""" - if not self.loaded: - self.load_model() - - # Check CUDA requirements if use_cuda is enabled - if use_cuda: - system_has_cuda = torch.cuda.is_available() - opencv_has_cuda = self.is_cuda_available() - - if not system_has_cuda: - raise RuntimeError("CUDA acceleration requested but CUDA is not available on your system. " - "Either install CUDA or disable CUDA acceleration in the node settings.") - - if not opencv_has_cuda: - raise RuntimeError("CUDA acceleration requested but your OpenCV build doesn't support CUDA. " - "Install a CUDA-enabled OpenCV build (via conda: 'conda install -c conda-forge opencv cudatoolkit') " - "or disable CUDA acceleration in the node settings.") - - # Convert to cv2 format if needed - cv2_image = self.convert_comfy_to_cv2(image) - - # Process the image - upscaled = self.model.upsample(cv2_image) - - # Convert back to original format - return self.convert_cv2_to_comfy(upscaled, image) - - def convert_comfy_to_cv2(self, tensor): - """Convert ComfyUI tensor (BHWC) to cv2 image (HWC)""" - # Handle both tensor and numpy array inputs - if isinstance(tensor, torch.Tensor): - # Ensure tensor is on CPU and convert to numpy - image = tensor.detach().cpu().numpy() - else: - image = tensor - - # Take first image if batch dimension exists - if len(image.shape) == 4: - image = image[0] - - # Convert float tensor [0,1] to uint8 [0,255] - if image.dtype == np.float32 or image.dtype == np.float64: - image = (image * 255).astype(np.uint8) - - # Convert RGB to BGR for OpenCV - if image.shape[2] == 3: # RGB - image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR) - - return image - - def convert_cv2_to_comfy(self, image, original_tensor=None): - """Convert cv2 image (HWC) back to ComfyUI tensor format (BHWC)""" - # Convert BGR to RGB - if image.shape[2] == 3: # BGR - image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) - - # Convert to float [0,1] if original was float - if original_tensor is not None and ( - isinstance(original_tensor, torch.Tensor) and original_tensor.dtype == torch.float32 or - isinstance(original_tensor, np.ndarray) and original_tensor.dtype == np.float32 - ): - image = image.astype(np.float32) / 255.0 - - # Add batch dimension if needed - if original_tensor is not None and len(original_tensor.shape) == 4: - image = np.expand_dims(image, axis=0) - - # Convert to torch tensor if original was tensor - if isinstance(original_tensor, torch.Tensor): - image = torch.from_numpy(image).to(original_tensor.device) - - return image - - def release(self): - """Release resources""" - self.model = None - self.loaded = False - - def __del__(self): - self.release() \ No newline at end of file diff --git a/src/super_resolution/models/test_opencv_sr.py b/src/super_resolution/models/test_opencv_sr.py deleted file mode 100644 index ce8563e..0000000 --- a/src/super_resolution/models/test_opencv_sr.py +++ /dev/null @@ -1,110 +0,0 @@ -import cv2 -import sys -import torch - -def test_opencv_version(): - print(f"OpenCV Version: {cv2.__version__}") - print(f"Python Version: {sys.version}") - print(f"CUDA available (torch): {torch.cuda.is_available()}") - if torch.cuda.is_available(): - print(f"CUDA version (torch): {torch.version.cuda}") - -def test_cuda_availability(): - print("\nTesting OpenCV CUDA Support:") - - # Check if CUDA module exists - print(f"cv2.cuda exists: {hasattr(cv2, 'cuda')}") - if hasattr(cv2, 'cuda'): - print(f"CUDA devices: {cv2.cuda.getCudaEnabledDeviceCount()}") - - # Check DNN CUDA backend availability - print("\nTesting DNN CUDA Backend:") - if hasattr(cv2.dnn, "getAvailableBackends"): - backends = cv2.dnn.getAvailableBackends() - print(f"Available backends: {backends}") - print(f"Has CUDA backend: {cv2.dnn.DNN_BACKEND_CUDA in backends}") - else: - print("cv2.dnn.getAvailableBackends not available") - - if hasattr(cv2.dnn, "getAvailableTargets"): - try: - targets = cv2.dnn.getAvailableTargets(cv2.dnn.DNN_BACKEND_CUDA) - print(f"Available CUDA targets: {targets}") - except Exception as e: - print(f"Error getting CUDA targets: {e}") - else: - print("cv2.dnn.getAvailableTargets not available") - -def test_dnn_superres_availability(): - print("\nTesting DNN SuperRes availability:") - print(f"cv2.dnn_superres exists: {hasattr(cv2, 'dnn_superres')}") - - if hasattr(cv2, 'dnn_superres'): - print("\nAvailable attributes in cv2.dnn_superres:") - for attr in dir(cv2.dnn_superres): - if not attr.startswith('__'): - print(f"- {attr}") - -def test_create_methods(): - print("\nTesting creation methods:") - - methods = [ - ("DnnSuperResImpl()", lambda: cv2.dnn_superres.DnnSuperResImpl()), - ("DnnSuperResImpl_create()", lambda: cv2.dnn_superres.DnnSuperResImpl_create()) - ] - - working_methods = [] - - for method_name, method_func in methods: - try: - model = method_func() - print(f"✓ {method_name} - SUCCESS") - working_methods.append((method_name, model)) - except Exception as e: - print(f"✗ {method_name} - FAILED: {str(e)}") - - return working_methods - -def test_model_methods(working_methods): - if not working_methods: - print("\nNo working creation methods found!") - return - - print("\nTesting model methods and CUDA setup on working instances:") - for method_name, model in working_methods: - print(f"\nTesting {method_name}:") - try: - print("Basic methods:") - print("- hasattr readModel:", hasattr(model, 'readModel')) - print("- hasattr setModel:", hasattr(model, 'setModel')) - print("- hasattr upsample:", hasattr(model, 'upsample')) - - print("\nCUDA-related methods:") - print("- hasattr setPreferableBackend:", hasattr(model, 'setPreferableBackend')) - if hasattr(model, 'setPreferableBackend'): - try: - model.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA) - print(" ✓ Successfully set CUDA backend") - except Exception as e: - print(f" ✗ Failed to set CUDA backend: {e}") - - print("- hasattr setPreferableTarget:", hasattr(model, 'setPreferableTarget')) - if hasattr(model, 'setPreferableTarget'): - try: - model.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA) - print(" ✓ Successfully set CUDA target") - except Exception as e: - print(f" ✗ Failed to set CUDA target: {e}") - - except Exception as e: - print(f"Error testing methods: {str(e)}") - -if __name__ == "__main__": - print("OpenCV Super Resolution CUDA Test Script") - print("=" * 40) - - test_opencv_version() - test_cuda_availability() - test_dnn_superres_availability() - working_methods = test_create_methods() - test_model_methods(working_methods) \ No newline at end of file diff --git a/src/super_resolution/nodes/__init__.py b/src/super_resolution/nodes/__init__.py deleted file mode 100644 index 6c3468d..0000000 --- a/src/super_resolution/nodes/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# Makes nodes directory a proper Python package \ No newline at end of file diff --git a/src/super_resolution/nodes/sr_nodes.py b/src/super_resolution/nodes/sr_nodes.py deleted file mode 100644 index 9b8c051..0000000 --- a/src/super_resolution/nodes/sr_nodes.py +++ /dev/null @@ -1,95 +0,0 @@ -class SuperResolutionModelLoader: - """Node that loads a super resolution model and outputs a reference to it""" - - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "model_type": (["FSRCNN", "FSRCNN-small", "EDSR", "ESPCN", "LapSRN", "VDSR"], { - "default": "FSRCNN", - "tooltip": "Select super resolution model type:\n" - "• FSRCNN: Good balance of quality and speed\n" - "• FSRCNN-small: Fastest, lower quality option\n" - "• ESPCN: Efficient for text and line art\n" - "• LapSRN: Better edge preservation\n" - "• EDSR: Highest quality, slower processing\n" - "• VDSR: Very Deep SR with sharp edge reconstruction\n" - }), - "scale_factor": (["2", "3", "4"], { - "default": "2", - "tooltip": "Upscaling factor:\n" - "• 2: Double resolution (e.g., 512x512 → 1024x1024)\n" - "• 3: Triple resolution (e.g., 512x512 → 1536x1536)\n" - "• 4: Quadruple resolution (e.g., 512x512 → 2048x2048)" - }), - } - } - - RETURN_TYPES = ("SR_MODEL",) - RETURN_NAMES = ("sr_model",) - FUNCTION = "load_model" - CATEGORY = "StreamPack/SuperResolution" - - def load_model(self, model_type, scale_factor): - """Load the selected super resolution model and return a reference""" - - # Import unified model class - from ..models import SuperResolutionModel - - # Import model download utilities - from ..utils.model_utils import download_model - - # On-demand download the model file if it doesn't exist - scale_key = f"x{scale_factor}" - model_path = None - - # Download the model if needed - model_path = download_model(model_type, scale_key) - if not model_path: - print(f"Warning: Could not download model {model_type} {scale_key}") - - # Create the model using our factory method - model = SuperResolutionModel.create(model_type, int(scale_factor)) - - # Load the model with the downloaded path if available - if model_path: - model.load_model(model_path) - else: - # Fall back to default paths - model.load_model() - - return (model,) - - -class SuperResolutionUpscale: - """Node that applies a super resolution model to an image""" - - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "image": ("IMAGE",), - "sr_model": ("SR_MODEL",), - "use_cuda": (["True", "False"], { - "default": "True", - "tooltip": "Whether to use CUDA for acceleration (requires CUDA-enabled OpenCV build)" - }), - } - } - - RETURN_TYPES = ("IMAGE",) - RETURN_NAMES = ("upscaled_image",) - FUNCTION = "upscale_image" - OUTPUT_NODE = True - CATEGORY = "StreamPack/SuperResolution" - - def upscale_image(self, image, sr_model, use_cuda): - """Upscale image using the provided super resolution model""" - - # Convert use_cuda string to boolean - use_cuda = use_cuda == "True" - - # Upscale the image with the use_cuda parameter - result = sr_model.upscale(image, use_cuda=use_cuda) - - return (result,) \ No newline at end of file diff --git a/src/super_resolution/utils/__init__.py b/src/super_resolution/utils/__init__.py deleted file mode 100644 index 849fa21..0000000 --- a/src/super_resolution/utils/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -# Utils module for ComfyUI SuperResolution -from .model_utils import download_model, get_model_path \ No newline at end of file diff --git a/src/super_resolution/utils/model_utils.py b/src/super_resolution/utils/model_utils.py deleted file mode 100644 index a0e109f..0000000 --- a/src/super_resolution/utils/model_utils.py +++ /dev/null @@ -1,149 +0,0 @@ -import os -import sys -import requests -import json -from pathlib import Path -import itertools -import shutil -import time - -# Fix the import path to use relative import -from ..models.sr_model import SuperResolutionModel - -def get_models_root_dir(): - """Get the root directory for downloaded SR models""" - # Find the ComfyUI_SuperResolution directory - base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) - # Create downloaded_models directory separate from code - return os.path.join(base_dir, "downloaded_models") - -def get_model_path(model_family, scale): - """Get the path to a model file, ensuring directories exist""" - # Normalize model family name to lowercase for consistency - model_family = model_family.lower() - - # Check if model family exists in our configurations - if model_family not in SuperResolutionModel.MODEL_CONFIGS: - print(f"Error: Unknown model family '{model_family}'") - return None - - # Get configuration for this model - config = SuperResolutionModel.MODEL_CONFIGS[model_family] - - # Convert scale to integer for comparison - scale_int = int(scale.replace('x', '')) - - # Check if scale is supported for this model - if scale_int not in config["supported_scales"]: - print(f"Error: Scale {scale} not supported for {model_family}") - return None - - # Get model filename from pattern - model_filename = config["filename_pattern"].format(scale=scale_int) - - # Create family-specific subdirectory - models_dir = os.path.join(get_models_root_dir(), model_family.upper()) - os.makedirs(models_dir, exist_ok=True) - - # Return the full path to the model file - return os.path.join(models_dir, model_filename) - -def download_model_file(url, destination): - """Download a model file with progress indication""" - # Skip if model already exists - if os.path.exists(destination): - print(f"Model already exists: {os.path.basename(destination)}") - return True - - print(f"Downloading {os.path.basename(destination)}...") - try: - # Download with progress indication for large files - with requests.get(url, stream=True) as r: - r.raise_for_status() - total_size = int(r.headers.get('content-length', 0)) - with open(destination, 'wb') as f: - if total_size > 1024*1024: # Show progress for files > 1MB - downloaded = 0 - start_time = time.time() - for chunk in r.iter_content(chunk_size=8192): - if chunk: - f.write(chunk) - downloaded += len(chunk) - # Show progress every 1MB - if downloaded % (1024*1024) < 8192: - percent = 100 * downloaded / total_size - elapsed = time.time() - start_time - speed = downloaded / (1024 * 1024 * elapsed) if elapsed > 0 else 0 - print(f" {percent:.1f}% ({downloaded/(1024*1024):.1f}MB) - {speed:.1f}MB/s") - else: - # For small files, just download without progress - shutil.copyfileobj(r.raw, f) - print(f"Downloaded {os.path.basename(destination)} successfully!") - return True - except Exception as e: - print(f"Error downloading {url}: {e}") - # Remove partial file if download failed - if os.path.exists(destination): - os.remove(destination) - return False - -def download_model(model_family, scale): - """Download a specific SR model on-demand if it doesn't exist""" - # Normalize model family name to lowercase for consistency - model_family_norm = model_family.lower() - - # Check if valid model family - if model_family_norm not in SuperResolutionModel.MODEL_CONFIGS: - print(f"Error: Unknown model family '{model_family}'") - return None - - config = SuperResolutionModel.MODEL_CONFIGS[model_family_norm] - - # Build model download URL map if we don't have a static one - model_urls = { - "fsrcnn": { - "base_url": "https://raw.githubusercontent.com/ryanontheinside/FSRCNN_Tensorflow/master/models/", - }, - "fsrcnn-small": { - "base_url": "https://raw.githubusercontent.com/ryanontheinside/FSRCNN_Tensorflow/master/models/", - }, - "edsr": { - "base_url": "https://raw.githubusercontent.com/ryanontheinside/EDSR_Tensorflow/master/models/", - }, - "espcn": { - "base_url": "https://raw.githubusercontent.com/ryanontheinside/TF-ESPCN/master/export/", - }, - "lapsrn": { - "base_url": "https://raw.githubusercontent.com/ryanontheinside/TF-LapSRN/master/export/", - }, - "vdsr": { - "base_url": "https://raw.githubusercontent.com/ryanontheinside/TF-VDSR/master/export/", - } - } - - # Check if we have a download URL for this model - if model_family_norm not in model_urls: - print(f"Error: No download URL for model family '{model_family}'") - return None - - # Get the model path - model_path = get_model_path(model_family, scale) - if not model_path: - return None - - # If model already exists, return its path - if os.path.exists(model_path): - return model_path - - # Extract the scale number from the format - scale_int = int(scale.replace('x', '')) - - # Get the base URL and create the download URL - base_url = model_urls[model_family_norm]["base_url"] - filename = config["filename_pattern"].format(scale=scale_int) - url = base_url + filename - - if download_model_file(url, model_path): - return model_path - else: - return None \ No newline at end of file