preprocessor.py

#  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)