How to load the specific finetuned model?

[JimReno]

Repect for the great commnity! Diffuers has already helped me a lot :). Currently I am using the StableDiffusionInstructPix2PixPipeline to train a instruct pix2pix model. I have trained the model and save many checkpoint files. The saved folder is shown as below:

The question is that How can I load the model in checkpoint-10000 or in checkpoint-15000? Is there any args I can pass to theStableDiffusionInstructPix2PixPipeline class? The raw inference code is:

...
model_id = "./uieb_raw_resolution512/" # <- replace this
pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(model_id, torch_dtype=torch.float16).to("cuda")
generator = torch.Generator("cuda").manual_seed(0)
...

which seems the StableDiffusionInstructPix2PixPipeline will automatically load the latest model when I pass the model id arg: ./uieb_raw_resolution512/.

[bmaxdk]

Can you try to specify your checkpoint in your saved path like this?

# Try to specify your checkpt
model_id = "./uieb_raw_resolution512/checkpoint-10000/"  # or 15000

pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(model_id, torch_dtype=torch.float16).to("cuda")
generator = torch.Generator("cuda").manual_seed(0)

[JimReno]

Thank you for the quick respoonse 👍 . I have run this code and show the error shown below:

  File "/root/autodl-tmp/diffusers/src/diffusers/configuration_utils.py", line 373, in load_config
    raise EnvironmentError(
OSError: Error no file named model_index.json found in directory ./uieb_raw_resolution512/checkpoint-10000/.

I guess the class StableDiffusionInstructPix2PixPipeline receive a parent folder uieb_raw_resolution512 where it contains model_index.jsonto find unet model, vae model and sampler. Solely give the checkpoint folder will not worl :( .
A hint may come from issue #9804 where the control net is defined by a checkpoint path:

controlnet = ControlNetModel.from_pretrained('./tdq_new_sdxl_controlnet_20241029/checkpoint-35000', torch_dtype = torch.float16)

vae = AutoencoderKL.from_pretrained('madebyollin/sdxl-vae-fp16-fix', torch_dtype = torch.float16)

pipeline = StableDiffusionXLControlNetPAGPipeline.from_pretrained(
    'stabilityai/stable-diffusion-xl-base-1.0',
    controlnet = controlnet,
    torch_dtype = torch.float16,
    # enable_pag = True,
    pag_applied_layers = ['mid'],
    vae = vae,
)

Based on that, I modified the code as below, it just doesn't throw any error. But I don't know if it is right:

from diffusers import StableDiffusionInstructPix2PixPipeline,UNet2DConditionModel

unet_model_id = "./uieb_raw_resolution512/checkpoint-10000/unet"
unet2dcon = UNet2DConditionModel.from_pretrained(unet_model_id, torch_dtype=torch.float16).to("cuda")
model_id = "./uieb_raw_resolution512/"
pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(model_id, torch_dtype=torch.float16).to("cuda")
pipe.unet = unet2dcon
generator = torch.Generator("cuda").manual_seed(0)

[bmaxdk]

It should be right approach. You load the unet model from your specific checkpoint directory and then assign the fine-tuned unet to the pipeline, replacing the original one. Due to this, you are loading the fine-tuned model for inference

[JimReno]

Thanks for the help :).