Help with fine tune Stable Diffusion v1-5 on Pytorchlightnig

🧨 Diffusers
1

For some reason when I try to initialize the optimizer and write
self.model.parameters()
I get an AttributeError:
'StableDiffusionPipeline' object has no attribute 'parameters'
This my model :

class StableDiffusionFineTuner(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.model_id = "runwayml/stable-diffusion-v1-5"
        self.model = StableDiffusionPipeline.from_pretrained(self.model_id, torch_dtype=torch.float16).to("cuda")
        self.loss = torch.nn.MSELoss()
        self.optimizer = torch.optim.Adam(self.model.parameters(),lr=1e-5,weight_decay=(1e-2,),betas=(0.9,0.999),eps=1e-08)

    def forward(self, input_text):
        return self.model(input_text).images[0]

    def training_step(self, batch, batch_idx):
        images, texts = batch
        outputs = self.forward(texts)
        loss = self.loss(outputs, images)
        self.log("train_loss", loss)
        return loss

    def validation_step(self, batch, batch_idx):
        images, texts = batch
        outputs = self.forward(texts)
        loss = self.loss(outputs, images)
        self.log("val_loss", loss)
        return loss

    def configure_optimizers(self):
        optimizer = self.optimizer
        return optimizer

I assume that this error is not related to my garbage code, but more of a system error. If anyone can point out the cause of this error and help me figure out what the problem is, I would be very grateful.

2

Hi,

The pipeline is only used at inference time. During fine-tuning, you need to load its individual building blocks (like the U-Net) and optimize them. It’s actually only the U-Net that gets fine-tuned, the VAE and text encoder are freezed as seen here.

Basically you would need to port this script to PyTorch Lightning: https://github.com/huggingface/diffusers/blob/main/examples/text_to_image/train_text_to_image.py.

1 Like
3

Oh wow thank you!