Optimizing text embedding raises issue of trying to backward through the graph second time
from io import BytesIO

import requests
import torch
from PIL import Image
from diffusers import DiffusionPipeline, DDIMScheduler
import numpy as np
import PIL
import torch
from packaging import version


if version.parse(version.parse(PIL.__version__).base_version) >= version.parse("9.1.0"):
    PIL_INTERPOLATION = {
        "linear": PIL.Image.Resampling.BILINEAR,
        "bilinear": PIL.Image.Resampling.BILINEAR,
        "bicubic": PIL.Image.Resampling.BICUBIC,
        "lanczos": PIL.Image.Resampling.LANCZOS,
        "nearest": PIL.Image.Resampling.NEAREST,
    }
else:
    PIL_INTERPOLATION = {
        "linear": PIL.Image.LINEAR,
        "bilinear": PIL.Image.BILINEAR,
        "bicubic": PIL.Image.BICUBIC,
        "lanczos": PIL.Image.LANCZOS,
        "nearest": PIL.Image.NEAREST,
    }
# ------------------------------------------------------------------------------

def preprocess(image):
    w, h = image.size
    w, h = (x - x % 32 for x in (w, h))  # resize to integer multiple of 32
    image = image.resize((w, h), resample=PIL_INTERPOLATION["lanczos"])
    image = np.array(image).astype(np.float32) / 255.0
    image = image[None].transpose(0, 3, 1, 2)
    image = torch.from_numpy(image)
    return 2.0 * image - 1.0


def get_text_embeddings(prompt, tokenizer, text_encoder, device):
    text_input = tokenizer(prompt,
                           padding="max_length",
                           max_length=tokenizer.model_max_length,
                           truncation=True,
                           return_tensors="pt", )

    text_embeddings_orig = text_encoder(text_input.input_ids.to(device))[0]

    return text_embeddings_orig

def get_latents(image, vae, generator, dtype, device):
    image = preprocess(image)
    image = image.to(device=device, dtype=dtype)
    init_latent_image_dist = vae.encode(image).latent_dist
    image_latents = init_latent_image_dist.sample(generator=generator)
    image_latents = 0.18215 * image_latents
    return image_latents

def get_pipeline_components(pipe):
    vae  = pipe.vae # the VAE model
    unet = pipe.unet # the U-Net model
    text_encoder = pipe.text_encoder # the text encoder model
    tokenizer = pipe.tokenizer # the tokenizer
    scheduler = pipe.scheduler # the scheduler
    return vae, unet, text_encoder, tokenizer, scheduler

def get_device():
    if torch.cuda.is_available():
        return torch.device("cuda")
    else:
        return torch.device("cpu")

# Set the device
device = get_device()

# Load the pipeline
pipe = DiffusionPipeline.from_pretrained(
    "CompVis/stable-diffusion-v1-4",
    safety_checker=None,
    torch_dtype=torch.float16,
).to(device)

pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)




# Load the image and text prompt
generator = torch.Generator("cuda").manual_seed(0)
seed = 0
prompt = "A photo of Barack Obama smiling with a big grin"
url = 'https://www.dropbox.com/s/6tlwzr73jd1r9yk/obama.png?dl=1'
response = requests.get(url)
init_image = Image.open(BytesIO(response.content)).convert("RGB")
init_image = init_image.resize((512, 512))

#image = pipe(prompt, generator=generator).images[0]



# Get the different components from the pipeline
vae, unet, text_encoder, tokenizer, scheduler = get_pipeline_components(pipe)
vae = vae.to(device).eval()
unet = unet.to(device).eval()
text_encoder = text_encoder.to(device).eval()

# Encode the text prompt and get the text embeddings
text_embeddings = get_text_embeddings(prompt, tokenizer, text_encoder, device=device)
text_embeddings = torch.nn.Parameter(text_embeddings, requires_grad=True)
text_embeddings_orig = text_embeddings.clone()

# Initialize the optimizer
optimizer = torch.optim.Adam(
    [text_embeddings],  # only optimize the embeddings
    lr=0.000001,
)


# Get Latent representation of the image
image_latents = get_latents(init_image, vae, generator, text_embeddings.dtype, device=device)

# Optimize the text embeddings to better reconstruct the init image
for _ in range(500):
    noise = torch.randn(image_latents.shape, dtype=text_embeddings.dtype).to(image_latents.device)
    timesteps = torch.randint(1000, (1,), device=image_latents.device)

    # Add noise to the latents according to the noise magnitude at each timestep
    # (this is the forward diffusion process)
    noisy_latents = scheduler.add_noise(image_latents, noise, timesteps)

    # Predict the noise residual
    noise_pred = unet(noisy_latents, timesteps, text_embeddings).sample

    loss = torch.nn.functional.mse_loss(noise_pred, noise, reduction="none").mean([1, 2, 3]).mean()

    loss.backward()

    optimizer.step()
    optimizer.zero_grad()

    print(loss.item())0.03802490234375╭───────────────────── Traceback (most recent call last) ──────────────────────╮│ /home/hashmat/Code/Projects/Diffusion_Attack/imagic/main.py:74 in    ││                                                                              ││   71 │                                                                       ││   72 │   loss = torch.nn.functional.mse_loss(noise_pred, noise, reduction="n ││   73 │                                                                       ││ ❱ 74 │   loss.backward()                                                     ││   75 │                                                                       ││   76 │   optimizer.step()                                                    ││   77 │   optimizer.zero_grad()                                               ││                                                                              ││ /home/hashmat/miniconda3/envs/pix2pix/lib/python3.11/site-packages/torch/_te ││ nsor.py:487 in backward                                                      ││                                                                              ││    484 │   │   │   │   create_graph=create_graph,                            ││    485 │   │   │   │   inputs=inputs,                                        ││    486 │   │   │   )                                                         ││ ❱  487 │   │   torch.autograd.backward(                                      ││    488 │   │   │   self, gradient, retain_graph, create_graph, inputs=inputs ││    489 │   │   )                                                             ││    490                                                                       ││                                                                              ││ /home/hashmat/miniconda3/envs/pix2pix/lib/python3.11/site-packages/torch/aut ││ ograd/init.py:200 in backward                                            ││                                                                              ││   197 │   # The reason we repeat same the comment below is that              ││   198 │   # some Python versions print out the first line of a multi-line fu ││   199 │   # calls in the traceback and some print out the last line          ││ ❱ 200 │   Variable.execution_engine.run_backward(  # Calls into the C++ eng ││   201 │   │   tensors, grad_tensors, retain_graph, create_graph, inputs,    ││   202 │   │   allow_unreachable=True, accumulate_grad=True)  # Calls into th ││   203                                                                        │╰──────────────────────────────────────────────────────────────────────────────╯RuntimeError: Trying to backward through the graph a second time (or directlyaccess saved tensors after they have already been freed). Saved intermediatevalues of the graph are freed when you call .backward() or autograd.grad().Specify retain_graph=True if you need to backward through the graph a secondtime or if you need to access saved tensors after calling backward.Process finished with exit code 1I am trying to optimize the text embedding in order to make the embedding output the given image. However, it raises the above error. However, I am not doing the backward pass on the loss twice. Can someone please explain the issue here?
2023-10-10T20:01:40Z
Hello, I’m doing something similar to your case (optimizing the text embedding) and I’m getting the same error. Did you manage to solve it?