Cannot load custom pipeline modules defined in the main script

Describe the bug

pipeline_util.register_modules tries to retrieve __module__ from pipeline modules and crashes for modules defined in the main class because the module __main__ does not contain a .

Reproduction

Create a pipeline with an own safetychecker class, e.g.:

class DummyChecker:
    def __call__(self, clip_input, images):
        images, False

create a pipe with that module:

pipe = StableDiffusionImg2ImgPipeline(
                                    vae,
                                    text_encoder,
                                    tokenizer,
                                    unet,
                                    scheduler,
                                    DummyChecker(),
                                    feature_extractor)

Logs

With an added `print(name, module.__module__)`:


vae diffusers.models.vae
text_encoder transformers.models.clip.modeling_clip
tokenizer transformers.models.clip.tokenization_clip
unet diffusers.models.unet_2d_condition
scheduler diffusers.schedulers.scheduling_lms_discrete
safety_checker __main__
(...)
  File "diffusers/src/diffusers/pipeline_utils.py", line 88, in register_modules
    pipeline_dir = module.__module__.split(".")[-2]
IndexError: list index out of range

### System Info

diffusers git checkout

[Karbo123]

Just found a simple solution, but not sure what the side effect is:

class PNDMPipeline(DiffusionPipeline):
    def __init__(self, unet, scheduler):
        super().__init__()
        scheduler = scheduler.set_format("pt")
        self.register_to_config(unet=("__main__", "Model"))
        self.register_to_config(scheduler=("diffusers", "PNDMScheduler"))
        self.unet = unet
        self.scheduler = scheduler

[allo-]

My script is using the Pipelines without safetychecker since weeks, I saw no point in hiding it. I think it’s 700 MB for something that only benefits one when one wants to provide a public service. shrugs

My current issue is, that I need to reload the model to change between generation/img2img/inpainting and I’d like to re-use the already loaded nets, so one can change the pipeline from a prompt in my script. And when I completely deleted the net in the old pipeline, I cannot pass it to the constructor of the new one. Constructing one to delete it again could also work, but I think when the old pipeline already filled the VRAM I cannot even construct a temporary checker. I didn’t put too much energy into trying to construct something to delete it again, though 😉

So … I could just plug it into my source checkout of diffusors, but I think each part of the pipeline will have the problem when someone wants to replace it by an own part, so maybe it is better to improve the utils class.