如何确保 .nonzero() 返回一个元素张量？

Question

[编辑以包含原始源代码]

我尝试运行在这里找到的代码：https://colab.research.google.com/drive/1roZqqhsdpCXZr8kgV_Bx_ABVBPgea3lX ?usp=sharing（链接自：https: //www.youtube.com/watch ?v=-lz30by8 -sU )

!pip install transformers diffusers lpips accelerate
from huggingface_hub import notebook_login
notebook_login()

import torch
from transformers import CLIPTextModel, CLIPTokenizer
from diffusers import AutoencoderKL, UNet2DConditionModel, LMSDiscreteScheduler
from tqdm.auto import tqdm
from torch import autocast
from PIL import Image
from matplotlib import pyplot as plt
import numpy
from torchvision import transforms as tfms

# For video display:
from IPython.display import HTML
from base64 import b64encode

# Set device
torch_device = "cuda" if torch.cuda.is_available() else "cpu"

# Load the autoencoder model which will be used to decode the latents into image space. 
vae = AutoencoderKL.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="vae", use_auth_token=True)

# Load the tokenizer and text encoder to tokenize and encode the text. 
tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14")
text_encoder = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14")

# The UNet model for generating the latents.
unet = UNet2DConditionModel.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="unet", use_auth_token=True)

# The noise scheduler
scheduler = LMSDiscreteScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000)

# To the GPU we go!
vae = vae.to(torch_device)
text_encoder = text_encoder.to(torch_device)
unet = unet.to(torch_device)

from google.colab import drive
drive.mount('/content/drive')

prompt = ["A digital illustration of a steampunk computer laboratory with clockwork machines, 4k, detailed, trending in artstation, fantasy vivid colors"]
height = 512
width = 768
num_inference_steps = 50
guidance_scale = 7.5
generator = torch.manual_seed(4)
batch_size = 1

# Prep text 
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"
)
with torch.no_grad():
  uncond_embeddings = text_encoder(uncond_input.input_ids.to(torch_device))[0] 
text_embeddings = torch.cat([uncond_embeddings, text_embeddings])

# Prep Scheduler
scheduler.set_timesteps(num_inference_steps)

# Prep latents
latents = torch.randn(
(batch_size, unet.in_channels, height // 8, width // 8),
generator=generator,
)

latents = latents.to(torch_device)
latents = latents * scheduler.sigmas[0] # Need to scale to match k

    # Loop
with autocast("cuda"):
    for i, t in tqdm(enumerate(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)
        sigma = scheduler.sigmas[i]
        latent_model_input = latent_model_input / ((sigma**2 + 1) ** 0.5)

        # predict the noise residual
        with torch.no_grad():
            noise_pred = 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 = scheduler.step(noise_pred, i, latents)["prev_sample"]

# scale and decode the image latents with vae
latents = 1 / 0.18215 * latents

with torch.no_grad():
    image = vae.decode(latents)

# Display
image = (image / 2 + 0.5).clamp(0, 1)
image = image.detach().cpu().permute(0, 2, 3, 1).numpy()
images = (image * 255).round().astype("uint8")
pil_images = [Image.fromarray(image) for image in images]
pil_images[0]

但是，我遇到了以下错误：

--------------------------------------------------------------------------- ValueError                                Traceback (most recent call last) <ipython-input-36-0fa46b18e9c1> in <module>
     48 
     49         # compute the previous noisy sample x_t -> x_t-1
---> 50         latents = scheduler.step(noise_pred, i, latents)["prev_sample"]
     51 
     52 # scale and decode the image latents with vae

/usr/local/lib/python3.8/dist-packages/diffusers/schedulers/scheduling_lms_discrete.py in step(self, model_output, timestep, sample, order, return_dict)
    216             timestep = timestep.to(self.timesteps.device)
    217 
--> 218         step_index = (self.timesteps == timestep).nonzero().item()
    219         sigma = self.sigmas[step_index]
    220 

ValueError: only one element tensors can be converted to Python scalars

此错误仅在循环的第二次迭代期间发生。第一次迭代顺利进行。

我打印了涉及的变量（noise_pred、i、latents）及其各自的维度。它们在第一次和第二次迭代期间具有相同的尺寸。

由于我在 Colab 上运行它，因此我无法直接访问 Scheduling_lms_discrete.py 中的底层代码

我该怎么做才能避免这个错误？这与 python 或 torch 的版本控制有关吗？（当前版本：python==3.8.16。torch==1.13.0+cu116）

谢谢！

Answer 1

尝试使用t而不是i参数scheduler.step：

latents = scheduler.step(noise_pred, t, latents)["prev_sample"]