mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-05-07 15:52:32 +08:00
c69a80d94e
- setup.py now works - Makes installation work a variety of ways, including making other packages dependent on this one for e.g. plugins - Fixes missing __init__.py issues - Fixes imports - Compatible with your existing scripts that rely on requirements.txt - Fixes error in comfy/ldm/models/diffusion/ddim.py - Fixes missing packages for other diffusers code in this repo
50 lines
1.6 KiB
Python
50 lines
1.6 KiB
Python
import os
|
|
from comfy_extras.chainner_models import model_loading
|
|
from comfy.sd import load_torch_file
|
|
from comfy import model_management
|
|
import torch
|
|
import comfy.utils
|
|
import folder_paths
|
|
|
|
class UpscaleModelLoader:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": { "model_name": (folder_paths.get_filename_list("upscale_models"), ),
|
|
}}
|
|
RETURN_TYPES = ("UPSCALE_MODEL",)
|
|
FUNCTION = "load_model"
|
|
|
|
CATEGORY = "loaders"
|
|
|
|
def load_model(self, model_name):
|
|
model_path = folder_paths.get_full_path("upscale_models", model_name)
|
|
sd = load_torch_file(model_path)
|
|
out = model_loading.load_state_dict(sd).eval()
|
|
return (out, )
|
|
|
|
|
|
class ImageUpscaleWithModel:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": { "upscale_model": ("UPSCALE_MODEL",),
|
|
"image": ("IMAGE",),
|
|
}}
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "upscale"
|
|
|
|
CATEGORY = "image/upscaling"
|
|
|
|
def upscale(self, upscale_model, image):
|
|
device = model_management.get_torch_device()
|
|
upscale_model.to(device)
|
|
in_img = image.movedim(-1,-3).to(device)
|
|
s = comfy.utils.tiled_scale(in_img, lambda a: upscale_model(a), tile_x=128 + 64, tile_y=128 + 64, overlap = 8, upscale_amount=upscale_model.scale)
|
|
upscale_model.cpu()
|
|
s = torch.clamp(s.movedim(-3,-1), min=0, max=1.0)
|
|
return (s,)
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"UpscaleModelLoader": UpscaleModelLoader,
|
|
"ImageUpscaleWithModel": ImageUpscaleWithModel
|
|
}
|