[Bug]: (TensorRT) Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0!

Checklist

• The issue exists after disabling all extensions
• The issue exists on a clean installation of webui
• The issue is caused by an extension, but I believe it is caused by a bug in the webui
• The issue exists in the current version of the webui
• The issue has not been reported before recently
• The issue has been reported before but has not been fixed yet

What happened?

When trying to use TensorRT to generate images, it fails with error "Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0! (when checking argument for argument mat1 in method wrapper_CUDA_addmm)"

Tbh, I'm not sure if it's an issue with the extension itself, or Forge, but I have friends who got it working in Forge, so I'm unsure what might be the issue here.

I got this issue before while trying to generate Export Engine, but now I get it when generating images only.

Steps to reproduce the problem

1. Open webui, try to generate image without SD Unet (None), it works normally
2. Select generated Unet (Ex:[TRT] mfcgPDXL_v10)
3. Unet model loads, but generation fails

What should have happened?

Unet loads, generates image, without using cpu device for any tensors

What browsers do you use to access the UI ?

Other

Sysinfo

sysinfo-2024-05-13-07-41.json

Console logs

Python 3.10.11 (tags/v3.10.11:7d4cc5a, Apr  5 2023, 00:38:17) [MSC v.1929 64 bit (AMD64)]
Version: f0.0.17v1.8.0rc-latest-276-g29be1da7
Commit hash: 29be1da7cf2b5dccfc70fbdd33eb35c56a31ffb7
removing nvidia-cudnn-cu11
Launching Web UI with arguments: --cuda-stream --disable-xformers
Total VRAM 12287 MB, total RAM 16314 MB
Set vram state to: NORMAL_VRAM
Device: cuda:0 NVIDIA GeForce RTX 3060 : native
Hint: your device supports --pin-shared-memory for potential speed improvements.
Hint: your device supports --cuda-malloc for potential speed improvements.
VAE dtype: torch.bfloat16
CUDA Stream Activated:  True
Using pytorch cross attention
ControlNet preprocessor location: E:\Programming\stable-diffusion-webui\models\ControlNetPreprocessor
Loading weights [529c72f6c3] from E:\Programming\stable-diffusion-webui\models\Stable-diffusion\mfcgPDXL_v10.safetensors
model_type EPS
UNet ADM Dimension 2816
2024-05-13 04:37:31,386 - ControlNet - INFO - ControlNet UI callback registered.
Running on local URL:  http://127.0.0.1:7860

To create a public link, set `share=True` in `launch()`.
Startup time: 46.5s (prepare environment: 11.9s, import torch: 9.5s, import gradio: 2.8s, setup paths: 3.1s, initialize shared: 0.3s, other imports: 1.8s, list SD models: 0.2s, load scripts: 13.5s, create ui: 2.8s, gradio launch: 0.2s).
Using pytorch attention in VAE
Working with z of shape (1, 4, 32, 32) = 4096 dimensions.
Using pytorch attention in VAE
extra {'cond_stage_model.clip_g.transformer.text_model.embeddings.position_ids', 'cond_stage_model.clip_l.text_projection', 'cond_stage_model.clip_l.logit_scale'}
loaded straight to GPU
To load target model SDXL
Begin to load 1 model
[Memory Management] Current Free GPU Memory (MB) =  6229.2353515625
[Memory Management] Model Memory (MB) =  0.03814697265625
[Memory Management] Minimal Inference Memory (MB) =  1024.0
[Memory Management] Estimated Remaining GPU Memory (MB) =  5205.197204589844
Moving model(s) has taken 0.04 seconds
Loading VAE weights specified in settings: E:\Programming\stable-diffusion-webui\models\VAE\sdxl_vae_fix.safetensors
To load target model SDXLClipModel
Begin to load 1 model
[Memory Management] Current Free GPU Memory (MB) =  6228.701171875
[Memory Management] Model Memory (MB) =  2144.3546981811523
[Memory Management] Minimal Inference Memory (MB) =  1024.0
[Memory Management] Estimated Remaining GPU Memory (MB) =  3060.3464736938477
Moving model(s) has taken 0.57 seconds
Model loaded in 52.9s (load weights from disk: 1.3s, forge load real models: 45.8s, load VAE: 4.1s, calculate empty prompt: 1.6s).
Activating unet: [TRT] mfcgPDXL_v10
Loading TensorRT engine: E:\Programming\stable-diffusion-webui\models\Unet-trt\mfcgPDXL_v10_871c8cac_cc86_sample=2x4x64x64-timesteps=2-encoder_hidden_states=2x77x2048-y=2x2816.trt

Loaded Profile: 0
        sample = [(2, 4, 64, 64), (2, 4, 64, 64), (2, 4, 64, 64)]
        timesteps = [(2,), (2,), (2,)]
        encoder_hidden_states = [(2, 77, 2048), (2, 77, 2048), (2, 77, 2048)]
        y = [(2, 2816), (2, 2816), (2, 2816)]
        latent = [(2, 4, 64, 64), (2, 4, 64, 64), (2, 4, 64, 64)]

  0%|                                                                                           | 0/20 [00:01<?, ?it/s]
Traceback (most recent call last):
  File "E:\Programming\stable-diffusion-webui\modules_forge\main_thread.py", line 37, in loop
    task.work()
  File "E:\Programming\stable-diffusion-webui\modules_forge\main_thread.py", line 26, in work
    self.result = self.func(*self.args, **self.kwargs)
  File "E:\Programming\stable-diffusion-webui\modules\txt2img.py", line 111, in txt2img_function
    processed = processing.process_images(p)
  File "E:\Programming\stable-diffusion-webui\modules\processing.py", line 752, in process_images
    res = process_images_inner(p)
  File "E:\Programming\stable-diffusion-webui\modules\processing.py", line 922, in process_images_inner
    samples_ddim = p.sample(conditioning=p.c, unconditional_conditioning=p.uc, seeds=p.seeds, subseeds=p.subseeds, subseed_strength=p.subseed_strength, prompts=p.prompts)
  File "E:\Programming\stable-diffusion-webui\modules\processing.py", line 1275, in sample
    samples = self.sampler.sample(self, x, conditioning, unconditional_conditioning, image_conditioning=self.txt2img_image_conditioning(x))
  File "E:\Programming\stable-diffusion-webui\modules\sd_samplers_kdiffusion.py", line 251, in sample
    samples = self.launch_sampling(steps, lambda: self.func(self.model_wrap_cfg, x, extra_args=self.sampler_extra_args, disable=False, callback=self.callback_state, **extra_params_kwargs))
  File "E:\Programming\stable-diffusion-webui\modules\sd_samplers_common.py", line 263, in launch_sampling
    return func()
  File "E:\Programming\stable-diffusion-webui\modules\sd_samplers_kdiffusion.py", line 251, in <lambda>
    samples = self.launch_sampling(steps, lambda: self.func(self.model_wrap_cfg, x, extra_args=self.sampler_extra_args, disable=False, callback=self.callback_state, **extra_params_kwargs))
  File "E:\Programming\stable-diffusion-webui\venv\lib\site-packages\torch\utils\_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "E:\Programming\stable-diffusion-webui\repositories\k-diffusion\k_diffusion\sampling.py", line 145, in sample_euler_ancestral
    denoised = model(x, sigmas[i] * s_in, **extra_args)
  File "E:\Programming\stable-diffusion-webui\venv\lib\site-packages\torch\nn\modules\module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "E:\Programming\stable-diffusion-webui\venv\lib\site-packages\torch\nn\modules\module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "E:\Programming\stable-diffusion-webui\modules\sd_samplers_cfg_denoiser.py", line 182, in forward
    denoised = forge_sampler.forge_sample(self, denoiser_params=denoiser_params,
  File "E:\Programming\stable-diffusion-webui\modules_forge\forge_sampler.py", line 88, in forge_sample
    denoised = sampling_function(model, x, timestep, uncond, cond, cond_scale, model_options, seed)
  File "E:\Programming\stable-diffusion-webui\ldm_patched\modules\samplers.py", line 289, in sampling_function
    cond_pred, uncond_pred = calc_cond_uncond_batch(model, cond, uncond_, x, timestep, model_options)
  File "E:\Programming\stable-diffusion-webui\ldm_patched\modules\samplers.py", line 258, in calc_cond_uncond_batch
    output = model.apply_model(input_x, timestep_, **c).chunk(batch_chunks)
  File "E:\Programming\stable-diffusion-webui\ldm_patched\modules\model_base.py", line 90, in apply_model
    model_output = self.diffusion_model(xc, t, context=context, control=control, transformer_options=transformer_options, **extra_conds).float()
  File "E:\Programming\stable-diffusion-webui\venv\lib\site-packages\torch\nn\modules\module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "E:\Programming\stable-diffusion-webui\venv\lib\site-packages\torch\nn\modules\module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "E:\Programming\stable-diffusion-webui\ldm_patched\ldm\modules\diffusionmodules\openaimodel.py", line 854, in forward
    emb = self.time_embed(t_emb)
  File "E:\Programming\stable-diffusion-webui\venv\lib\site-packages\torch\nn\modules\module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "E:\Programming\stable-diffusion-webui\venv\lib\site-packages\torch\nn\modules\module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "E:\Programming\stable-diffusion-webui\venv\lib\site-packages\torch\nn\modules\container.py", line 215, in forward
    input = module(input)
  File "E:\Programming\stable-diffusion-webui\venv\lib\site-packages\torch\nn\modules\module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "E:\Programming\stable-diffusion-webui\venv\lib\site-packages\torch\nn\modules\module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "E:\Programming\stable-diffusion-webui\ldm_patched\modules\ops.py", line 98, in forward
    return super().forward(*args, **kwargs)
  File "E:\Programming\stable-diffusion-webui\venv\lib\site-packages\torch\nn\modules\linear.py", line 114, in forward
    return F.linear(input, self.weight, self.bias)
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0! (when checking argument for argument mat1 in method wrapper_CUDA_addmm)
Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0! (when checking argument for argument mat1 in method wrapper_CUDA_addmm)
*** Error completing request
*** Arguments: ('task(74lerspdo1fd51g)', <gradio.routes.Request object at 0x0000026E18879630>, '', '', [], 20, 'Euler a', 1, 1, 7, 512, 512, False, 0.7, 2, 'Latent', 0, 0, 0, 'Use same checkpoint', 'Use same sampler', '', '', [], 0, False, '', 0.8, -1, False, -1, 0, 0, 0, 0.0, 4, 512, 512, True, 'None', 'None', 0, True, 'keyword prompt', 'keyword1, keyword2', 'None', 'textual inversion first', 'None', '0.7', 'None', ControlNetUnit(input_mode=<InputMode.SIMPLE: 'simple'>, use_preview_as_input=False, batch_image_dir='', batch_mask_dir='', batch_input_gallery=[], batch_mask_gallery=[], generated_image=None, mask_image=None, hr_option='Both', enabled=False, module='None', model='None', weight=1, image=None, resize_mode='Crop and Resize', processor_res=512, threshold_a=64, threshold_b=64, guidance_start=0, guidance_end=1, pixel_perfect=False, control_mode='Balanced', save_detected_map=True), ControlNetUnit(input_mode=<InputMode.SIMPLE: 'simple'>, use_preview_as_input=False, batch_image_dir='', batch_mask_dir='', batch_input_gallery=[], batch_mask_gallery=[], generated_image=None, mask_image=None, hr_option='Both', enabled=False, module='None', model='None', weight=1, image=None, resize_mode='Crop and Resize', processor_res=512, threshold_a=64, threshold_b=64, guidance_start=0, guidance_end=1, pixel_perfect=False, control_mode='Balanced', save_detected_map=True), ControlNetUnit(input_mode=<InputMode.SIMPLE: 'simple'>, use_preview_as_input=False, batch_image_dir='', batch_mask_dir='', batch_input_gallery=[], batch_mask_gallery=[], generated_image=None, mask_image=None, hr_option='Both', enabled=False, module='None', model='None', weight=1, image=None, resize_mode='Crop and Resize', processor_res=512, threshold_a=64, threshold_b=64, guidance_start=0, guidance_end=1, pixel_perfect=False, control_mode='Balanced', save_detected_map=True), False, 7, 1, 'Constant', 0, 'Constant', 0, 1, 'enable', 'MEAN', 'AD', 1, False, 1.01, 1.02, 0.99, 0.95, False, 0.5, 2, False, 256, 2, 0, False, False, 3, 2, 0, 0.35, True, 'bicubic', 'bicubic', False, 0, 'anisotropic', 0, 'reinhard', 100, 0, 'subtract', 0, 0, 'gaussian', 'add', 0, 100, 127, 0, 'hard_clamp', 5, 0, 'None', 'None', False, 'MultiDiffusion', 768, 768, 64, 4, False, False, False, False, False, 'positive', 'comma', 0, False, False, 'start', '', 1, '', [], 0, '', [], 0, '', [], True, False, False, False, False, False, False, 0, False) {}
    Traceback (most recent call last):
      File "E:\Programming\stable-diffusion-webui\modules\call_queue.py", line 57, in f
        res = list(func(*args, **kwargs))
    TypeError: 'NoneType' object is not iterable

---

Additional information

-I updated GPU drives today
-Tried with and without --xformers/--disable-xformers and --always-gpu
-In TensorRT, had to comment out @swap_sdpa decorator to build engine succesfuly
-"Other" browser is Opera

Reproduce the failure by selecting the generated TensorRT UNet and starting image generation. Start with modules_forge/forge_sampler.py and follow the traceback through ldm_patched/modules/samplers.py, model_base.py, and ldm/modules/diffusionmodules/openaimodel.py. Done means the selected TensorRT UNet generates an image without a CPU/CUDA device mismatch.