Unverified Commit 8cfd677c authored by Joey Ballentine's avatar Joey Ballentine Committed by GitHub
Browse files

Replace chainner_models with Spandrel package (#2146)

* Replace chainner_models with Spandrel

* Update to latest spandrel

* Use spandrel_foss instead

* update spandrel to new FOSS-compliant version
parent ffc4b7c3
This diff is collapsed.
This diff is collapsed.
This diff is collapsed.
This diff is collapsed.
import os import os
from comfy_extras.chainner_models import model_loading from spandrel import ModelLoader, ImageModelDescriptor
from comfy import model_management from comfy import model_management
import torch import torch
import comfy.utils import comfy.utils
...@@ -20,7 +20,11 @@ class UpscaleModelLoader: ...@@ -20,7 +20,11 @@ class UpscaleModelLoader:
sd = comfy.utils.load_torch_file(model_path, safe_load=True) sd = comfy.utils.load_torch_file(model_path, safe_load=True)
if "module.layers.0.residual_group.blocks.0.norm1.weight" in sd: if "module.layers.0.residual_group.blocks.0.norm1.weight" in sd:
sd = comfy.utils.state_dict_prefix_replace(sd, {"module.":""}) sd = comfy.utils.state_dict_prefix_replace(sd, {"module.":""})
out = model_loading.load_state_dict(sd).eval() out = ModelLoader().load_from_state_dict(sd).eval()
if not isinstance(out, ImageModelDescriptor):
raise Exception("Upscale model must be a single-image model.")
return (out, ) return (out, )
...@@ -61,7 +65,7 @@ class ImageUpscaleWithModel: ...@@ -61,7 +65,7 @@ class ImageUpscaleWithModel:
if tile < 128: if tile < 128:
raise e raise e
upscale_model.cpu() upscale_model.to("cpu")
s = torch.clamp(s.movedim(-3,-1), min=0, max=1.0) s = torch.clamp(s.movedim(-3,-1), min=0, max=1.0)
return (s,) return (s,)
......
This diff is collapsed.
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment