[Bug]: Model doesn't load on pytorch nightly 2.6

[Panchovix]

Describe the bug

Hi, the error consists that, if using the .pt YOLO models with adetailer on newer pytorch versions (nightly), it won't load because unsafe loading.

I have fixed by modifying code to load the file.

From

@contextmanager
def change_torch_load():
    orig = torch.load
    try:
        torch.load = safe.unsafe_torch_load
        yield
    finally:
        torch.load = orig

To

@contextmanager
def change_torch_load():
    orig = torch.load
    try:
        # Add the DetectionModel to safe globals
        torch.serialization.add_safe_globals(['ultralytics.nn.tasks.DetectionModel'])
        
        # Create a custom loader function
        def custom_load(*args, **kwargs):
            # Force weights_only to False if not explicitly set
            if 'weights_only' not in kwargs:
                kwargs['weights_only'] = False
            return orig(*args, **kwargs)
        
        # Replace torch.load with our custom loader
        torch.load = custom_load
        yield
    finally:
        torch.load = orig

Steps to reproduce

1. Install torch nightly 2.6.0 from November or newer.
2. Tried to use adetailer
3. Get error

Screenshots

No response

Console logs, from start to end.

*** Error running postprocess_image: G:\Stable difussion\stable-diffusion-webui-reForge\extensions\adetailer\scripts\!adetailer.py
    Traceback (most recent call last):
      File "G:\Stable difussion\stable-diffusion-webui-reForge\modules\scripts.py", line 931, in postprocess_image
        script.postprocess_image(p, pp, *script_args)
      File "G:\Stable difussion\stable-diffusion-webui-reForge\extensions\adetailer\aaaaaa\traceback.py", line 173, in wrapper
        raise error from None
    ValueError:
    ╭──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╮
    │                                                     System info                                                      │
    │ ┏━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓ │
    │ ┃             ┃ Value                                                                                              ┃ │
    │ ┡━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩ │
    │ │    Platform │ Windows-10-10.0.26100-SP0                                                                          │ │
    │ │      Python │ 3.10.11 (tags/v3.10.11:7d4cc5a, Apr  5 2023, 00:38:17) [MSC v.1929 64 bit (AMD64)]                 │ │
    │ │     Version │ f1.3.6dev2-v1.10.1RC-previous-1276-g77a4f0ce                                                       │ │
    │ │      Commit │ 77a4f0ce8d9a6bdfa33f58e9b42a00f541bd3538                                                           │ │
    │ │ Commandline │ ['launch.py', '--xformers', '--skip-python-version-check', '--always-gpu', '--listen', '--api',    │ │
    │ │             │ '--enable-insecure-extension-access', '--administrator', '--port', '7980', '--max-batch-count',    │ │
    │ │             │ '32', '--lyco-dir', 'K:\\LoRA_cosas\\Lora', '--lora-dir', 'K:\\LoRA_cosas\\Lora', '--gradio-auth', │ │
    │ │             │ '--disable-nan-check', '--disable-tls-verify', '--cuda-malloc',              │ │
    │ │             │ '--cuda-stream', '--pin-shared-memory']                                                            │ │
    │ │   Libraries │ {'torch': '2.6.0.dev20241117+cu124', 'torchvision': '0.20.0.dev20241118+cu124', 'ultralytics':     │ │
    │ │             │ '8.3.29', 'mediapipe': '0.10.14'}                                                                  │ │
    │ └─────────────┴────────────────────────────────────────────────────────────────────────────────────────────────────┘ │
    │                                                        Inputs                                                        │
    │ ┏━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓ │
    │ ┃                 ┃ Value                                                                                          ┃ │
    │ ┡━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩ │
    │ │          prompt │ 1girl,   │ │
    │ │                 │ masterpiece, best quality, absurdres                                                           │ │
    │ │ negative_prompt │ lowres, worst quality, low quality                             │ │
    │ │          n_iter │ 4                                                                                              │ │
    │ │      batch_size │ 1                                                                                              │ │
    │ │           width │ 896                                                                                            │ │
    │ │          height │ 1152                                                                                           │ │
    │ │    sampler_name │ DPM++ SDE                                                                                      │ │
    │ │       enable_hr │ True                                                                                           │ │
    │ │     hr_upscaler │ 2x-AnimeSharpV3                                                                                │ │
    │ │      checkpoint │ normals\xlmodels\illu-noob\illustriousXLPersonalMerge_v30Noob10based.safetensors               │ │
    │ │             vae │ sdxl_vae_fixedfp16.safetensors                                                                 │ │
    │ │            unet │ Automatic                                                                                      │ │
    │ └─────────────────┴────────────────────────────────────────────────────────────────────────────────────────────────┘ │
    │                 ADetailer                                                                                            │
    │ ┏━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━┓                                                                            │
    │ ┃                     ┃ Value           ┃                                                                            │
    │ ┡━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━┩                                                                            │
    │ │             version │ 24.11.1         │                                                                            │
    │ │            ad_model │ face_yolov8n.pt │                                                                            │
    │ │           ad_prompt │                 │                                                                            │
    │ │  ad_negative_prompt │                 │                                                                            │
    │ │ ad_controlnet_model │ None            │                                                                            │
    │ │              is_api │ False           │                                                                            │
    │ └─────────────────────┴─────────────────┘                                                                            │
    │ ╭─────────────────────────────── Traceback (most recent call last) ────────────────────────────────╮                 │
    │ │ G:\Stable difussion\stable-diffusion-webui-reForge\modules_forge\patch_basic.py:13 in loader     │                 │
    │ │                                                                                                  │                 │
    │ │   12 │   │   try:                                                                                │                 │
    │ │ ❱ 13 │   │   │   result = original_loader(*args, **kwargs)                                       │                 │
    │ │   14 │   │   except Exception as e:                                                              │                 │
    │ │                                                                                                  │                 │
    │ │ G:\Stable                                                                                        │                 │
    │ │ difussion\stable-diffusion-webui-reForge\venv\lib\site-packages\torch\serialization.py:1450 in   │                 │
    │ │ load                                                                                             │                 │
    │ │                                                                                                  │                 │
    │ │   1449 │   │   │   │   │   except pickle.UnpicklingError as e:                                   │                 │
    │ │ ❱ 1450 │   │   │   │   │   │   raise pickle.UnpicklingError(_get_wo_message(str(e))) from None   │                 │
    │ │   1451 │   │   │   │   return _load(                                                             │                 │
    │ ╰──────────────────────────────────────────────────────────────────────────────────────────────────╯                 │
    │ UnpicklingError: Weights only load failed. This file can still be loaded, to do so you have two options, do those │
    │ steps only if you trust the source of the checkpoint.                                                             │
    │         (1) Re-running `torch.load` with `weights_only` set to `False` will likely succeed, but it can result in     │
    │ arbitrary code execution. Do it only if you got the file from a trusted source.                                      │
    │         (2) Alternatively, to load with `weights_only=True` please check the recommended steps in the following      │
    │ error message.                                                                                                       │
    │         WeightsUnpickler error: Unsupported global: GLOBAL ultralytics.nn.tasks.DetectionModel was not an allowed    │
    │ global by default. Please use `torch.serialization.add_safe_globals([DetectionModel])` or the                        │
    │ `torch.serialization.safe_globals([DetectionModel])` context manager to allowlist this global if you trust this      │
    │ class/function.                                                                                                      │
    │                                                                                                                      │
    │ Check the documentation of torch.load to learn more about types accepted by default with weights_only                │
    │ https://pytorch.org/docs/stable/generated/torch.load.html.                                                           │
    │                                                                                                                      │
    │ During handling of the above exception, another exception occurred:                                                  │
    │                                                                                                                      │
    │ ╭─────────────────────────────── Traceback (most recent call last) ────────────────────────────────╮                 │
    │ │ G:\Stable difussion\stable-diffusion-webui-reForge\extensions\adetailer\aaaaaa\traceback.py:153  │                 │
    │ │ in wrapper                                                                                       │                 │
    │ │                                                                                                  │                 │
    │ │   152 │   │   try:                                                                               │                 │
    │ │ ❱ 153 │   │   │   return func(*args, **kwargs)                                                   │                 │
    │ │   154 │   │   except Exception as e:                                                             │                 │
    │ │                                                                                                  │                 │
    │ │ G:\Stable                                                                                        │                 │
    │ │ difussion\stable-diffusion-webui-reForge\extensions\adetailer\scripts\!adetailer.py:916 in       │                 │
    │ │ postprocess_image                                                                                │                 │
    │ │                                                                                                  │                 │
    │ │    915 │   │   │   │   │   continue                                                              │                 │
    │ │ ❱  916 │   │   │   │   is_processed |= self._postprocess_image_inner(p, pp, args, n=n)           │                 │
    │ │    917                                                                                           │                 │
    │ │                                                                                                  │                 │
    │ │ G:\Stable                                                                                        │                 │
    │ │ difussion\stable-diffusion-webui-reForge\extensions\adetailer\scripts\!adetailer.py:830 in       │                 │
    │ │ _postprocess_image_inner                                                                         │                 │
    │ │                                                                                                  │                 │
    │ │    829 │   │   │   │   ad_model = self.get_ad_model(args.ad_model)                               │                 │
    │ │ ❱  830 │   │   │   │   pred = ultralytics_predict(                                               │                 │
    │ │    831 │   │   │   │   │   ad_model,                                                             │                 │
    │ │                                                                                                  │                 │
    │ │ G:\Stable                                                                                        │                 │
    │ │ difussion\stable-diffusion-webui-reForge\extensions\adetailer\adetailer\ultralytics.py:27 in     │                 │
    │ │ ultralytics_predict                                                                              │                 │
    │ │                                                                                                  │                 │
    │ │   26 │                                                                                           │                 │
    │ │ ❱ 27 │   model = YOLO(model_path)                                                                │                 │
    │ │   28 │   apply_classes(model, model_path, classes)                                               │                 │
    │ │                                                                                                  │                 │
    │ │ G:\Stable                                                                                        │                 │
    │ │ difussion\stable-diffusion-webui-reForge\venv\lib\site-packages\ultralytics\models\yolo\model.py │                 │
    │ │ :23 in __init__                                                                                  │                 │
    │ │                                                                                                  │                 │
    │ │    22 │   │   │   # Continue with default YOLO initialization                                    │                 │
    │ │ ❱  23 │   │   │   super().__init__(model=model, task=task, verbose=verbose)                      │                 │
    │ │    24                                                                                            │                 │
    │ │                                                                                                  │                 │
    │ │ G:\Stable                                                                                        │                 │
    │ │ difussion\stable-diffusion-webui-reForge\venv\lib\site-packages\ultralytics\engine\model.py:145  │                 │
    │ │ in __init__                                                                                      │                 │
    │ │                                                                                                  │                 │
    │ │    144 │   │   else:                                                                             │                 │
    │ │ ❱  145 │   │   │   self._load(model, task=task)                                                  │                 │
    │ │    146                                                                                           │                 │
    │ │                                                                                                  │                 │
    │ │ G:\Stable                                                                                        │                 │
    │ │ difussion\stable-diffusion-webui-reForge\venv\lib\site-packages\ultralytics\engine\model.py:285  │                 │
    │ │ in _load                                                                                         │                 │
    │ │                                                                                                  │                 │
    │ │    284 │   │   if Path(weights).suffix == ".pt":                                                 │                 │
    │ │ ❱  285 │   │   │   self.model, self.ckpt = attempt_load_one_weight(weights)                      │                 │
    │ │    286 │   │   │   self.task = self.model.args["task"]                                           │                 │
    │ │                                                                                                  │                 │
    │ │ G:\Stable                                                                                        │                 │
    │ │ difussion\stable-diffusion-webui-reForge\venv\lib\site-packages\ultralytics\nn\tasks.py:910 in   │                 │
    │ │ attempt_load_one_weight                                                                          │                 │
    │ │                                                                                                  │                 │
    │ │    909 │   """Loads a single model weights."""                                                   │                 │
    │ │ ❱  910 │   ckpt, weight = torch_safe_load(weight)  # load ckpt                                   │                 │
    │ │    911 │   args = {**DEFAULT_CFG_DICT, **(ckpt.get("train_args", {}))}  # combine model and def  │                 │
    │ │                                                                                                  │                 │
    │ │ G:\Stable                                                                                        │                 │
    │ │ difussion\stable-diffusion-webui-reForge\venv\lib\site-packages\ultralytics\nn\tasks.py:837 in   │                 │
    │ │ torch_safe_load                                                                                  │                 │
    │ │                                                                                                  │                 │
    │ │    836 │   │   │   else:                                                                         │                 │
    │ │ ❱  837 │   │   │   │   ckpt = torch.load(file, map_location="cpu")                               │                 │
    │ │    838                                                                                           │                 │
    │ │                                                                                                  │                 │
    │ │ G:\Stable difussion\stable-diffusion-webui-reForge\modules_forge\patch_basic.py:29 in loader     │                 │
    │ │                                                                                                  │                 │
    │ │   28 │   │   │   │   │   │   exp += f'You may try again now and Forge will download models ag    │                 │
    │ │ ❱ 29 │   │   │   raise ValueError(exp)                                                           │                 │
    │ │   30 │   │   return result                                                                       │                 │
    │ ╰──────────────────────────────────────────────────────────────────────────────────────────────────╯                 │
    │ ValueError: Weights only load failed. This file can still be loaded, to do so you have two options, do those      │
    │ steps only if you trust the source of the checkpoint.                                                             │
    │         (1) Re-running `torch.load` with `weights_only` set to `False` will likely succeed, but it can result in     │
    │ arbitrary code execution. Do it only if you got the file from a trusted source.                                      │
    │         (2) Alternatively, to load with `weights_only=True` please check the recommended steps in the following      │
    │ error message.                                                                                                       │
    │         WeightsUnpickler error: Unsupported global: GLOBAL ultralytics.nn.tasks.DetectionModel was not an allowed    │
    │ global by default. Please use `torch.serialization.add_safe_globals([DetectionModel])` or the                        │
    │ `torch.serialization.safe_globals([DetectionModel])` context manager to allowlist this global if you trust this      │
    │ class/function.                                                                                                      │
    │                                                                                                                      │
    │ Check the documentation of torch.load to learn more about types accepted by default with weights_only                │
    │ https://pytorch.org/docs/stable/generated/torch.load.html.                                                           │
    │ File corrupted: G:\Stable                                                                                            │
    │ difussion\stable-diffusion-webui-reForge\models\diffusers\models--Bingsu--adetailer\snapshots\b0a075fd35454c86bb453a │
    │ 1ca06b29ffee704c20\face_yolov8n.pt                                                                                   │
    │ Forge has tried to move the corrupted file to G:\Stable                                                              │
    │ difussion\stable-diffusion-webui-reForge\models\diffusers\models--Bingsu--adetailer\snapshots\b0a075fd35454c86bb453a │
    │ 1ca06b29ffee704c20\face_yolov8n.pt.corrupted                                                                         │
    │ You may try again now and Forge will download models again.                                                          │
    │                                                                                                                      │
    ╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯

List of installed extensions

No response

[dosubot]

@Panchovix ¡Hola! Soy Dosu y estoy aquí para ayudarte a solucionar problemas, responder preguntas y convertirte en un colaborador. Estoy aquí para ayudarte mientras esperas a un mantenedor humano.

The issue you're encountering is related to the way PyTorch nightly 2.6 handles model loading, which has become stricter with its safe loading mechanism. The error suggests that the YOLO model's DetectionModel class is not in the list of allowed globals, which is causing the torch.load function to fail when trying to load the model.

The modification you made to the change_torch_load function is a good approach to address this issue. Here's a breakdown of what your changes do:

1. Add Safe Globals: By using torch.serialization.add_safe_globals(['ultralytics.nn.tasks.DetectionModel']), you're allowing the DetectionModel class to be safely unpickled, which is necessary for loading the YOLO model.

2. Custom Loader Function: The custom loader function ensures that the weights_only parameter is set to False by default, which allows the full model to be loaded rather than just the weights. This is important because the error message indicates that setting weights_only to False can resolve the issue, provided the source is trusted.

These changes should help in loading the YOLO models with the newer PyTorch versions. Just ensure that you trust the source of the model files, as setting weights_only to False can execute arbitrary code during the loading process.

To continue talking to Dosu, mention @dosu.

[Bing-su]

I updated the model file and verified that it is loaded in 2.6 nightly.

I'll try to think of another way to do this rather than using the weights_only keyword for now.