From 82b7b817a34ad9021137e1f08de287440f48eb03 Mon Sep 17 00:00:00 2001
From: City <125218114+city96@users.noreply.github.com>
Date: Fri, 15 Sep 2023 19:14:37 +0200
Subject: [PATCH] Local model support

---
 README.md                |  8 +++++++-
 comfy_latent_upscaler.py | 19 +++++++++++++++----
 2 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/README.md b/README.md
index f836f43..3da8199 100644
--- a/README.md
+++ b/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.
diff --git a/comfy_latent_upscaler.py b/comfy_latent_upscaler.py
index 081ad1c..047542b 100644
--- a/comfy_latent_upscaler.py
+++ b/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)
-		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 = 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
 		)
-		# weights = f"./latent-upscaler-v{model.version}_SD{latent_ver}-x{scale_factor}.safetensors"
+
+		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=filename)
+			)
 
 		model.load_state_dict(load_file(weights))
 		lt = samples["samples"]