Issue with .to("cuda") on Space and ZeroGPU - RuntimeError: Expected all tensors to be on the same device

[thesab]

Hi everyone,

I’m encountering an issue when running my code on Space using ZeroGPU. I’ve set .to("cuda") on the pipeline, but I can’t figure out what’s going wrong. Locally, with two A100 GPUs, everything works fine.

Here is the complete error I’m getting:

Traceback (most recent call last):  
  File "/usr/local/lib/python3.10/site-packages/spaces/zero/wrappers.py", line 256, in thread_wrapper  
    res = future.result()  
  File "/usr/local/lib/python3.10/concurrent/futures/_base.py", line 451, in result  
    return self.__get_result()  
  File "/usr/local/lib/python3.10/concurrent/futures/_base.py", line 403, in __get_result  
    raise self._exception  
  File "/usr/local/lib/python3.10/concurrent/futures/thread.py", line 58, in run  
    result = self.fn(*self.args, *kwargs)  
  File "/home/user/app/app.py", line 49, in infer  
    image = pipe(  
  File "/usr/local/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context  
    return func(*args, **kwargs)  
  File "/usr/local/lib/python3.10/site-packages/diffusers/pipelines/flux/pipeline_flux.py", line 684, in __call__  
    latents, latent_image_ids = self.prepare_latents(  
  File "/usr/local/lib/python3.10/site-packages/diffusers/pipelines/flux/pipeline_flux.py", line 522, in prepare_latents  
    latent_image_ids = self._prepare_latent_image_ids(batch_size, height // 2, width // 2, device, dtype)  
  File "/usr/local/lib/python3.10/site-packages/diffusers/pipelines/flux/pipeline_flux.py", line 431, in _prepare_latent_image_ids  
    latent_image_ids[..., 1] = latent_image_ids[..., 1] + torch.arange(height)[:, None]  
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!  

It seems the issue is due to tensors being on different devices, in this case cuda:0 and cpu. I’ve checked my code and I believe I’ve transferred everything to cuda, so I’m not sure what I’m missing.

Has anyone experienced similar issues or have any suggestions on how to resolve this? Any help would be greatly appreciated!

Thanks in advance!

[John6666]

I’m guessing since I don’t have the code, but in Zero GPU space, CUDA is assumed to be absent from within a function with decorators or from the global scope. And .to(“cuda”) in a function without decorators has no effect.
The other is forgetting to do .to(“cuda”) in unobtrusive tensors such as inputs.

[thesab]

Hi John, this is the code:

if torch.cuda.is_available():
    torch_dtype = torch.bfloat16
else:
    torch_dtype = torch.float32

# Load Flux
pipe = FluxPipeline.from_pretrained(MODEL_PRE_TRAINED_ID, torch_dtype=torch.bfloat16).to(DEVICE)

# Load fine-tuned model
pipe.load_lora_weights(MODEL_FINE_TUNED_ID, adapter_name="default", weight_name="outfit-generator.safetensors")


@spaces.GPU(duration=65)
def infer(
    prompt,
    negative_prompt="",
    style = DEFAULT_STYLE_NAME,
    seed=42,
    randomize_seed=False,
    width=512,
    height=512,
    guidance_scale=4.5,
    num_inference_steps=40,
    progress=gr.Progress(track_tqdm=True),
):

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    generator = torch.Generator(device=DEVICE).manual_seed(seed)
    #print("STYLE SELECTED: ", style)
    prompt, negative_prompt = apply_style(style, prompt)
    #print(" PROMPT : ", prompt)

    # Esegui l'inferenza
    image = pipe(
            prompt=prompt,
            num_images_per_prompt=1,
            guidance_scale=guidance_scale,
            num_inference_steps=num_inference_steps,
            width=width,
            height=height,
            generator=generator,
        ).images[0]

    return image, seed

[John6666]

LoRA, or rather PEFT, can also cause that error. Comment it out and try it. It’s midnight here, so good night.

[thesab]

Thank you so much, John. I will follow your advice. Good night, I don’t want to keep you up.

[John6666]

Something, your code was almost the same as mine and there were no errors, but is there any part that is different between my code and yours (probably the dependency)?
It is offloaded from the Zero GPU space, but the following is the space where it actually worked.

test flux lora - a Hugging Face Space by John6666

[thesab]

It seems that something might have been corrupted in the Spaces, as I had to recreate the repository and now, “magically,” it works.

[thesab]

Outfit Generator - a Hugging Face Space by thesab

[John6666]

It’s that bug again… The cause is unknown, but it sometimes occurs and can be fixed by recreating the space. It might be an undiscovered hardware defect or something.
As long as it works, it’s OK.

[thesab]

As long as it works, it’s all good!

[mlwong]

Had the same error: ZeroGPU issues: Assertion `srcIndex < srcSelectDimSize` failed

After redeploying does the URL remain the same? I can’t casually change it and would rather have only the CPU version if the URL will change.

[John6666]

As it is a new creation, I think it will change. As a workaround, if you change the name of the currently existing space and create a new space in its place, you can achieve the same URL and keep the previous space without deleting it.