Re-train the openpose ControlNet model with DWPose annotated data

[huchenlei]

Hi, Thanks for your great work!

One thing I noticed previously in ControlNet is that the openpose model is not performing well on Controlling hands. Even very accuate pose is provided (Through manual editing). I was suspecting that was caused by training set of the ControlNet model containing large amount of mis-annotated data.

Is there a existing work/plan to re-train the openpose ControlNet model with DWPose annotaed data?

Thanks for your work again!

[yzd-v]

Yes, we also notice the situation. Later we will consider training a new ControlNet model for DWPose.

[huchenlei]

I am working on porting the preprocessor to ControlNet extension here: https://github.com/Mikubill/sd-webui-controlnet/pull/1856

Somehow if I run the test I get following error:

  File "D:\stable-diffusion-webui\extensions\sd-webui-controlnet\tests\annotator_tests\openpose_tests\openpose_e2e_test.py", line 17, in setUp
    self.detector.load_model()
  File "D:\stable-diffusion-webui\extensions\sd-webui-controlnet\annotator\openpose\__init__.py", line 207, in load_model
    self.dw_pose_estimation = Wholebody(dw_modelpath, device=self.device)
  File "D:\stable-diffusion-webui\extensions\sd-webui-controlnet\annotator\openpose\wholebody.py", line 37, in __init__
    self.pose_estimator = init_pose_estimator(pose_config, pose_ckpt, device=device)
  File "D:\stable-diffusion-webui\venv\lib\site-packages\mmpose\apis\inference.py", line 109, in init_model
    ckpt = load_checkpoint(model, checkpoint, map_location='cpu')
  File "D:\stable-diffusion-webui\venv\lib\site-packages\mmengine\runner\checkpoint.py", line 639, in load_checkpoint
    raise RuntimeError(
RuntimeError: No state_dict found in checkpoint file D:\stable-diffusion-webui\extensions\sd-webui-controlnet\annotator\downloads\openpose\dw-ll_ucoco_384.pth

The same pth file works fine if I run it in the Gradio demo code, or just run the constructor alone. I am not very familiar with torch. Did I set anything wrong that could cause this?

[yzd-v]

I follow your code and try openpose_e2e_test.py. But I do not meet the problem. The problem is here when the point is None. After I fix this, I get the final output as this:

[yzd-v]

Maybe the pth is broken? Or there are some problems with the file dir?

[huchenlei]

Ok I found the cause. It was A1111 banning loading unsafe .pth file:

class RestrictedUnpickler(pickle.Unpickler):
    extra_handler = None

    def persistent_load(self, saved_id):
        assert saved_id[0] == 'storage'

        try:
            return TypedStorage(_internal=True)
        except TypeError:
            return TypedStorage()  # PyTorch before 2.0 does not have the _internal argument

    def find_class(self, module, name):
        if self.extra_handler is not None:
            res = self.extra_handler(module, name)
            if res is not None:
                return res

        if module == 'collections' and name == 'OrderedDict':
            return getattr(collections, name)
        if module == 'torch._utils' and name in ['_rebuild_tensor_v2', '_rebuild_parameter', '_rebuild_device_tensor_from_numpy']:
            return getattr(torch._utils, name)
        if module == 'torch' and name in ['FloatStorage', 'HalfStorage', 'IntStorage', 'LongStorage', 'DoubleStorage', 'ByteStorage', 'float32', 'BFloat16Storage']:
            return getattr(torch, name)
        if module == 'torch.nn.modules.container' and name in ['ParameterDict']:
            return getattr(torch.nn.modules.container, name)
        if module == 'numpy.core.multiarray' and name in ['scalar', '_reconstruct']:
            return getattr(numpy.core.multiarray, name)
        if module == 'numpy' and name in ['dtype', 'ndarray']:
            return getattr(numpy, name)
        if module == '_codecs' and name == 'encode':
            return encode
        if module == "pytorch_lightning.callbacks" and name == 'model_checkpoint':
            import pytorch_lightning.callbacks
            return pytorch_lightning.callbacks.model_checkpoint
        if module == "pytorch_lightning.callbacks.model_checkpoint" and name == 'ModelCheckpoint':
            import pytorch_lightning.callbacks.model_checkpoint
            return pytorch_lightning.callbacks.model_checkpoint.ModelCheckpoint
        if module == "__builtin__" and name == 'set':
            return set

        # Forbid everything else.
        raise Exception(f"global '{module}/{name}' is forbidden")

Currently the workaround seems to be pass --disable-safe-unpickle to A1111. I am going to discuss with A1111 Devs on this. Thanks for your response!

[continue-revolution]

@huchenlei use this https://github.com/continue-revolution/sd-webui-segment-anything/blob/master/scripts/sam.py#L79-L83

[huchenlei]

@huchenlei use this https://github.com/continue-revolution/sd-webui-segment-anything/blob/master/scripts/sam.py#L79-L83

Thanks for pointer! I discussed with AUTO on discord. Issue already resolved.

[yzd-v]

@huchenlei We will release a new version with onnx. It will not rely on any mmxx.

[p0mad]

@huchenlei @yzd-v @continue-revolution Can you guys please guide me through installing or using this in controlnet iniside A1111? I'm not using controlnet-nighty as recommanded here

I have controlnet installaed inside A1111 and i just downloaded the pth model and put it inside extensions/sd-webui-controlnet/models

and it does not included the yaml file!

can i use the model now? should i set the resolution to 384x288 as model trainde on it?

Thanks Best regards

[huchenlei]

@p0mad This repo is not an A1111 extension. The .pth file is also not an ControlNet model so should not be placed in extensions/sd-webui-controlnet/models.

See https://github.com/Mikubill/sd-webui-controlnet/discussions/1863 for more details on how to use it in A1111 extension. The main branch is rolled back as lvmin does not want to introduce cpp dependency.

Also please submit a new issue instead of continuing the thread here, as this issue is tracking the re-train ControlNet openpose model task.

[yzd-v]

@huchenlei Now we can avoid installing mmcv through onnxruntime. Could it be available for A1111?

[fckbauer]

Yes, we also notice the situation. Later we will consider training a new ControlNet model for DWPose.

what happened to this? any updates? Thanks! I realized we have the same problem. (me and thread OP)

[fckbauer]

@huchenlei use this https://github.com/continue-revolution/sd-webui-segment-anything/blob/master/scripts/sam.py#L79-L83

Thanks for pointer! I discussed with AUTO on discord. Issue already resolved

How was it resolved? I want to know since aI am having the same problem. Thank you! 🙏🏻

[werran2]

Yes, we also notice the situation. Later we will consider training a new ControlNet model for DWPose.

looking forward to seeing the model soon

[fckbauer]

where is the new model? thanks

[tanghengjian]

hi guys, i test the controlnet openpose will not accurately produce similar number persons, such as the origial pose has three persons, while the output inference has four persons. So i want to re-train the openpose model with multi persons cases. Could the loss only care about the image's pose rather not the original full image pixel info? please @huchenlei @huchenlei @werran2 @yzd-v @fckbauer take a look!