Gradient for text embedding input to

emei: <p>I am having an issue to compute gradients after using the Stable diffusion model from hugging face at the input text embeddings. I use StableDiffusionPipeline from hugging face and use the pretrained model : generator = StableDiffusionPipeline.from_pretrained(“runwayml/stable-diffusion-v1-5”). I have text and use the text embeddings (from CLIP tokenizer and CLIPTextModel both from huggingface) as input to the generator. I have a loss, for example the loss is some function of the generated image (lets say i set the loss to sum of all values in the top half of the image and minimize it as close to zero). I want to update the text embedding by.backpropagation of the gradients.However I am unable to obtain gradients for the text embeddng even after requires_grad=True. I get an error which says “The .grad attribute of a Tensor that is not a leaf Tensor is being accessed. Its .grad attribute won’t be populated during autograd.backward(). If you indeed want the .grad field to be populated for a non-leaf Tensor, use .retain_grad() on the non-leaf Tensor. If you access the non-leaf Tensor by mistake, make sure you access the leaf Tensor instead.” I tried a number of different things but somehow doesnt seem to work. Someone know what the problem could be or a solution for this. thanks.</p>

ptrblck: <p>The error message points towards the error trying to access the <code>.grad</code> attribute of a non-leaf tensor as seen in this example:</p>
<pre><code class="lang-python">lin1 = nn.Linear(1, 1)
lin2 = nn.Linear(1, 1)

x = torch.randn(1, 1)
intermediate = lin1(x)
out = lin2(intermediate)

out.mean().backward()
# try to access .grad of a non-leaf
print(intermediate.is_leaf)
# False
print(intermediate.grad)
# None
# UserWarning: The .grad attribute of a Tensor that is not a leaf Tensor is being accessed. Its .grad attribute won't be populated during autograd.backward(). If you indeed want the .grad field to be populated for a non-leaf Tensor, use .retain_grad() on the non-leaf Tensor. If you access the non-leaf Tensor by mistake, make sure you access the leaf Tensor instead. See github.com/pytorch/pytorch/pull/30531 for more informations. (Triggered internally at aten/src/ATen/core/TensorBody.h:489.)

# fix it as suggested in the warning
intermediate = lin1(x)
intermediate.retain_grad()
out = lin2(intermediate)
out.mean().backward()
print(intermediate.grad)
# tensor([[0.9695]])
</code></pre>

emei: <p>I had tried retain.grad() and it didnt work. Below, I provide a small working example of something similar to what i want to do (most of the code taken from the huggingface example). I am missing something, but still not able to figure it out :</p>
<pre><code class="lang-auto">import torch
from transformers import CLIPTextModel, CLIPTokenizer
from diffusers import AutoencoderKL, UNet2DConditionModel, PNDMScheduler
from diffusers import LMSDiscreteScheduler
from tqdm.auto import tqdm

vae = AutoencoderKL.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="vae")
tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14")
text_encoder = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14")
unet = UNet2DConditionModel.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="unet")
scheduler = LMSDiscreteScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000)

torch_device = "cuda"
vae.to(torch_device)
text_encoder.to(torch_device)
unet.to(torch_device)

prompt = ["a photograph of an astronaut"]

height = 512                        # default height of Stable Diffusion
width = 512                         # default width of Stable Diffusion
num_inference_steps = 100           # Number of denoising steps
guidance_scale = 7.5                # Scale for classifier-free guidance
generator = torch.manual_seed(0)    # Seed generator to create the inital latent noise
learning_rate = 0.1
batch_size = len(prompt)

text_input = tokenizer(prompt, padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt")
text_embeddings = text_encoder(text_input.input_ids.to(torch_device))[0]
text_embeddings.retain_grad()

latents = torch.randn(
    (batch_size, unet.in_channels, height // 8, width // 8),
    generator=generator,
)
latents = latents.to(torch_device)
scheduler.set_timesteps(num_inference_steps)
latents = latents * scheduler.init_noise_sigma

scheduler.set_timesteps(num_inference_steps)
unet.train()

for t in tqdm(scheduler.timesteps):
    latent_model_input = scheduler.scale_model_input(latents, timestep=t)

    # predict the noise residual
    with torch.no_grad():
        noise_pred = unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample

    # compute the previous noisy sample x_t -&gt; x_t-1
    latents = scheduler.step(noise_pred, t, latents).prev_sample
    image = vae.decode(latents).sample
    loss = sum(sum(sum(sum(image))))
    loss.backward()
    print(text_embeddings.grad)
    # None
    print(text_embeddings.requires_grad)
    # True
</code></pre>

ptrblck: <p><code>text_embeddings</code> is used in a <code>no_grad</code> context so why would you expect to see gradients in this tensor?</p>

emei: <p>Yes , thats right. i was playing around with the code a lot (initially using the stablediffusionpipeline vs using the unet) so somehow that line ended up not being commented. I overlooked that one, thanks !</p>