Get the image and the gradient at the same time

[Zoybzo]

I am new to stable diffusion. I am trying to train a llm with stable diffusion, my loss need the output image to calculate. At first I use the pipeline to generate the image, but I found that there is no gradient, then I found that example file in this
file https://github.com/huggingface/diffusers/tree/main/examples/text_to_image , but I cannot get the image with it. I don't need to train the unet, is there a way to get the image and the gradient?

[tolgacangoz]

In the pipelines, there is a decorator torch method on top of their __call__ and several other methods. It avoids preserving gradients while inferencing. Could you examine the examples in these docs pages? In these examples, you can run the overall pipeline by separating its components and eliminating the aforementioned torch method in order to preserve gradients.

[Zoybzo]

Thanks for your reply! I still have some problems. There is the test. I added with torch.enable_grad(): before the generator, but there was sill no loss. As a comparision, I deleted the generator part, and got the loss.

class CustomModel(nn.Module):
    def __init__(self):
        super(CustomModel, self).__init__()
        # Stable diffusion
        self.generator = StableDiffusionPipeline.from_pretrained(
            "runwayml/stable-diffusion-v1-5",
            local_files_only=True,
        )
        self.tokenizer = CLIPTokenizer.from_pretrained(
            "openai/clip-vit-large-patch14",
            use_fast=True,
        )
        self.text_encoder = CLIPTextModel.from_pretrained(
            "openai/clip-vit-large-patch14", local_files_only=True
        )
        # custom layer
        self.linear = nn.Linear(768, 768)

    def forward(self, x):
        x = self.tokenizer(
            x,
            return_tensors="pt",
            padding="max_length",
            truncation=True,
            max_length=77,
        )
        x = self.text_encoder(**x).last_hidden_state
        x = self.linear(x)
        # with torch.enable_grad():
        #     x = self.generator(
        #         prompt_embeds=x,
        #         num_inference_steps=50,
        #     )
        # x = x.images[0]  # PIL.Image
        # convert to tensor
        # x = transforms.ToTensor()(x)
        return x


# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# load the model
logger.info("Loading the model")
model = CustomModel()
# generate the input
prompt = ["A beautiful landscape with mountains and rivers"]
# run the model
logger.info("Running the model")
output = model(prompt)
# logger the type of the output
logger.info("Output type: {}", type(output))
logger.info("Output shape: {}", output.shape)
# generate the target
# target = torch.randn(3, 512, 512)
target = torch.randn(1, 77, 768)
# calculate the loss
loss = nn.MSELoss()(output, target)
# print the loss
print(loss)
loss.backward()

The result for it:

And this is the result for the code with generator:

Do I need to define the components separated? Does the code correct?

[tolgacangoz]

Don't use StableDiffusionPipeline, use its U-Net and VAE separately as in the examples I shared. You need to get rid of StableDiffusionPipeline's torch.no_grad(). I think we might not need with torch.enable_grad(): here.

[Zoybzo]

Thank you! Now I got the gradient and the image. My code is there:

model_name = "runwayml/stable-diffusion-v1-5"


class CustomModel(nn.Module):
    def __init__(self, model_name, num_inference_steps, device):
        super(CustomModel, self).__init__()
        self.device = device
        self.unet = UNet2DConditionModel.from_pretrained(
            model_name, subfolder="unet", use_safetensors=True, local_files_only=True
        ).to(self.device)
        self.scheduler = UniPCMultistepScheduler.from_pretrained(
            model_name, subfolder="scheduler", local_files_only=True
        )
        self.scheduler.set_timesteps(num_inference_steps)
        self.linear = nn.Linear(64, 64)
        self.unet.requires_grad_(True)

    def forward(
        self,
        text_embeddings,
        batch_size,
        height,
        width,
        generator,
        torch_device,
    ):
        latents = torch.randn(
            (batch_size, self.unet.config.in_channels, height // 8, width // 8),
            generator=generator,
            device=torch_device,
        )
        latents = latents * self.scheduler.init_noise_sigma
        for t in tqdm(self.scheduler.timesteps):
            # expand the latents if we are doing classifier-free guidance to avoid doing two forward passes.
            latent_model_input = torch.cat([latents] * 2)
            latent_model_input = self.scheduler.scale_model_input(
                latent_model_input, timestep=t
            )
            # predict the noise residual
            with torch.no_grad():
                noise_pred = self.unet(
                    latent_model_input, t, encoder_hidden_states=text_embeddings
                ).sample
            # perform guidance
            noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
            noise_pred = noise_pred_uncond + guidance_scale * (
                noise_pred_text - noise_pred_uncond
            )
            # compute the previous noisy sample x_t -> x_t-1
            latents = self.scheduler.step(noise_pred, t, latents).prev_sample
        # scale and decode the image latents with vae
        latents = 1 / 0.18215 * latents
        logger.info(f"Latents shape: {latents.shape}")
        latents = self.linear(latents)
        return latents
        # with torch.no_grad():


# param
torch_device = "cuda"
prompt = ["a photograph of an astronaut riding a horse"]
height = 512  # default height of Stable Diffusion
width = 512  # default width of Stable Diffusion
num_inference_steps = 3  # Number of denoising steps
guidance_scale = 7.5  # Scale for classifier-free guidance
generator = torch.Generator(device=torch_device)
batch_size = len(prompt)
# load models
logger.info("Loading models")
vae = AutoencoderKL.from_pretrained(
    model_name, subfolder="vae", use_safetensors=True, local_files_only=True
)
tokenizer = CLIPTokenizer.from_pretrained(model_name, subfolder="tokenizer")
text_encoder = CLIPTextModel.from_pretrained(
    model_name, subfolder="text_encoder", use_safetensors=True, local_files_only=True
)
model = CustomModel(model_name, num_inference_steps, torch_device)
vae.to(torch_device)
text_encoder.to(torch_device)
model.to(torch_device)
# input
logger.info("Generating input")
text_input = tokenizer(
    prompt,
    padding="max_length",
    max_length=tokenizer.model_max_length,
    truncation=True,
    return_tensors="pt",
)

with torch.no_grad():
    text_embeddings = text_encoder(text_input.input_ids.to(torch_device))[0]
max_length = text_input.input_ids.shape[-1]
uncond_input = tokenizer(
    [""] * batch_size, padding="max_length", max_length=max_length, return_tensors="pt"
)
uncond_embeddings = text_encoder(uncond_input.input_ids.to(torch_device))[0]
text_embeddings = torch.cat([uncond_embeddings, text_embeddings])
# model
logger.info("Generating image")
latents = model(
    text_embeddings,
    batch_size,
    height,
    width,
    generator,
    torch_device,
)
image = vae.decode(latents).sample
image = (image / 2 + 0.5).clamp(0, 1).squeeze()
image = image.permute(0, 1, 2) * 255
logger.info(f"Image type: {type(image)}")
logger.info(f"Image shape: {image.shape}")
# generate the target
target = torch.randn_like(image)
loss = torch.nn.MSELoss()(image, target)
print(loss)
loss.backward()
logger.info(f"Loss: {loss}")

And the result:

But I still have a question. There is with torch.no_grad() before using unet in the CustomModel.forward in the code. Why the gradient cound be backward with the torch.no_grad()?

[tolgacangoz]

AFAIK, the components having with torch.no_grad() don't join gradient calculations. If you want them to join, then remove their with torch.no_grad(). Your vae doesn't have with torch.no_grad(), so it can join gradient calculations.