1
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-08-02 23:14:49 +08:00

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
This commit is contained in:
Joey Ballentine
2024-05-26 13:44:17 -04:00
committed by GitHub
parent ffc4b7c30e
commit 8cfd677cc0
53 changed files with 8 additions and 17722 deletions

View File

@@ -1,5 +1,5 @@
import os
from comfy_extras.chainner_models import model_loading
from spandrel import ModelLoader, ImageModelDescriptor
from comfy import model_management
import torch
import comfy.utils
@@ -20,7 +20,11 @@ class UpscaleModelLoader:
sd = comfy.utils.load_torch_file(model_path, safe_load=True)
if "module.layers.0.residual_group.blocks.0.norm1.weight" in sd:
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, )
@@ -61,7 +65,7 @@ class ImageUpscaleWithModel:
if tile < 128:
raise e
upscale_model.cpu()
upscale_model.to("cpu")
s = torch.clamp(s.movedim(-3,-1), min=0, max=1.0)
return (s,)