ComfyUI/comfy/ldm/data/util.py
Benjamin Berman c69a80d94e Create a setup.py and automatically select the correct pytorch binaries for the current platform and supported devices
- 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
2023-03-29 14:30:37 -07:00

24 lines
635 B
Python

import torch
from comfy.ldm.modules.midas.api import load_midas_transform
class AddMiDaS(object):
def __init__(self, model_type):
super().__init__()
self.transform = load_midas_transform(model_type)
def pt2np(self, x):
x = ((x + 1.0) * .5).detach().cpu().numpy()
return x
def np2pt(self, x):
x = torch.from_numpy(x) * 2 - 1.
return x
def __call__(self, sample):
# sample['jpg'] is tensor hwc in [-1, 1] at this point
x = self.pt2np(sample['jpg'])
x = self.transform({"image": x})["image"]
sample['midas_in'] = x
return sample