Spaces:
Build error
Build error
File size: 2,251 Bytes
26575f5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 |
# preprocessor.py
import torch
from PIL import Image
from controlnet_aux_local import NormalBaeDetector # Assuming this local package is available
class Preprocessor:
MODEL_ID = "lllyasviel/Annotators"
def __init__(self):
self.model = None
self.name = ""
def load(self, name: str) -> None:
if name == self.name:
return
elif name == "NormalBae":
print("Loading NormalBae")
device = "cuda" if torch.cuda.is_available() else "cpu"
# The .to(device) call should add the .device attribute
self.model = NormalBaeDetector.from_pretrained(self.MODEL_ID).to(device)
if torch.cuda.is_available():
torch.cuda.empty_cache()
self.name = name
else:
raise ValueError
return
def __call__(self, image: Image.Image, **kwargs) -> Image.Image:
device = "cuda" if torch.cuda.is_available() else "cpu"
# Check if the model has the 'device' attribute before accessing its 'type'
if hasattr(self.model, 'device'):
if self.model.device.type != device:
print(f"Moving preprocessor model to {device}") # Debug print
# Attempt to move the model if it's not on the correct device
try:
self.model.to(device)
except Exception as e:
print(f"Error moving preprocessor model to {device}: {e}")
pass # Continue and let the next line potentially raise an error
else:
# If .device attribute is missing, assume it's not on the correct device and try to move it
# This might happen if NormalBaeDetector is not a standard torch.nn.Module
print("Warning: Preprocessor model has no .device attribute. Attempting to move to correct device.")
try:
self.model.to(device) # Attempt to move even if .device is missing
except Exception as e:
print(f"Error attempting to move preprocessor model without .device attribute: {e}")
pass # Continue and let the next line potentially raise an error
return self.model(image, **kwargs)
|