Local model support

README.md

@@ -9,7 +9,7 @@ Very similar to my [latent interposer](https://github.com/city96/SD-Latent-Inter
 
 ### ComfyUI
 
-To install it, simply clone this repo to your custom_nodes folder using the following command: git clone https://github.com/city96/SD-Latent-Upscaler custom_nodes/SD-Latent-Upscaler.
+To install it, simply clone this repo to your custom_nodes folder using the following command: `git clone https://github.com/city96/SD-Latent-Upscaler custom_nodes/SD-Latent-Upscaler`.
 
 Alternatively, you can download the [comfy_latent_upscaler.py](https://github.com/city96/SD-Latent-Upscaler/blob/main/comfy_latent_upscaler.py) file to your ComfyUI/custom_nodes folder as well. You may need to install hfhub using the command pip install huggingface-hub inside your venv.
 
@@ -19,6 +19,12 @@ If you need the model weights for something else, they are [hosted on HF](https:
 
 Currently not supported but it should be possible to use it at the hires-fix part.
 
+### Local models
+
+The node pulls the required files from huggingface hub by default. You can create a `models` folder and place the modules there if you have a flaky connection or prefer to use it completely offline, it will load them locally instead. The path should be: `ComfyUI/custom_nodes/SD-Latent-Upscaler/models`
+
+Alternatively, just clone the entire HF repo to it: `git clone https://huggingface.co/city96/SD-Latent-Upscaler custom_nodes/SD-Latent-Upscaler/models`
+
 ### Usage/principle
 
 Usage is fairly simple. You use it anywhere where you would upscale a latent. If you need a higher scale factor (e.g. x4), simply chain two of the upscalers.

comfy_latent_upscaler.py

@@ -1,3 +1,4 @@
+import os
 import torch
 import torch.nn as nn
 from safetensors.torch import load_file
@@ -69,11 +70,21 @@ class LatentUpscaler:
 
 	def upscale(self, samples, latent_ver, scale_factor):
 		model = Upscaler(scale_factor)
+		filename = f"latent-upscaler-v{model.version}_SD{latent_ver}-x{scale_factor}.safetensors"
+		local = os.path.join(
+			os.path.join(os.path.dirname(os.path.realpath(__file__)),"models"),
+			filename
+		)
+
+		if os.path.isfile(local):
+			print("LatentUpscaler: Using local model")
+			weights = local
+		else:
+			print("LatentUpscaler: Using HF Hub model")
 			weights = str(hf_hub_download(
 				repo_id="city96/SD-Latent-Upscaler",
-			filename=f"latent-upscaler-v{model.version}_SD{latent_ver}-x{scale_factor}.safetensors")
+				filename=filename)
 			)
-		# weights = f"./latent-upscaler-v{model.version}_SD{latent_ver}-x{scale_factor}.safetensors"
 
 		model.load_state_dict(load_file(weights))
 		lt = samples["samples"]