Can't use any trt that isn't 512x512 batch 1

[Marco Tundo]

Anytime I try and and build a .trt model that isn't 512x512 batch 1 from an .onnx, the .trt is created without issues, but upon inference, I get:

  0%|                                                                                           | 0/20 [00:00<?, ?it/s]
*** Error completing request
*** Arguments: ('task(n26ahbotll6ugbl)', 'bird', '', [], 20, 0, False, False, 1, 1, 7, -1.0, -1.0, 0, 0, 0, False, 768, 768, False, 0.7, 2, 'Latent', 0, 0, 0, 0, '', '', [], <gradio.routes.Request object at 0x0000023B9EB9D960>, 0, False, False, 'positive', 'comma', 0, False, False, '', 1, '', [], 0, '', [], 0, '', [], True, False, False, False, 0) {}
    Traceback (most recent call last):
      File "C:\sd\stable-diffusion-webui\modules\call_queue.py", line 58, in f
        res = list(func(*args, **kwargs))
      File "C:\sd\stable-diffusion-webui\modules\call_queue.py", line 37, in f
        res = func(*args, **kwargs)
      File "C:\sd\stable-diffusion-webui\modules\txt2img.py", line 62, in txt2img
        processed = processing.process_images(p)
      File "C:\sd\stable-diffusion-webui\modules\processing.py", line 627, in process_images
        res = process_images_inner(p)
      File "C:\sd\stable-diffusion-webui\modules\processing.py", line 746, 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 "C:\sd\stable-diffusion-webui\modules\processing.py", line 999, in sample
        samples = self.sampler.sample(self, x, conditioning, unconditional_conditioning, image_conditioning=self.txt2img_image_conditioning(x))
      File "C:\sd\stable-diffusion-webui\modules\sd_samplers_kdiffusion.py", line 464, in sample
        samples = self.launch_sampling(steps, lambda: self.func(self.model_wrap_cfg, x, extra_args={
      File "C:\sd\stable-diffusion-webui\modules\sd_samplers_kdiffusion.py", line 303, in launch_sampling
        return func()
      File "C:\sd\stable-diffusion-webui\modules\sd_samplers_kdiffusion.py", line 464, in <lambda>
        samples = self.launch_sampling(steps, lambda: self.func(self.model_wrap_cfg, x, extra_args={
      File "C:\sd\stable-diffusion-webui\venv\lib\site-packages\torch\utils\_contextlib.py", line 115, in decorate_context
        return func(*args, **kwargs)
      File "C:\sd\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 "C:\sd\stable-diffusion-webui\venv\lib\site-packages\torch\nn\modules\module.py", line 1501, in _call_impl
        return forward_call(*args, **kwargs)
      File "C:\sd\stable-diffusion-webui\modules\sd_samplers_kdiffusion.py", line 183, in forward
        x_out = self.inner_model(x_in, sigma_in, cond=make_condition_dict(cond_in, image_cond_in))
      File "C:\sd\stable-diffusion-webui\venv\lib\site-packages\torch\nn\modules\module.py", line 1501, in _call_impl
        return forward_call(*args, **kwargs)
      File "C:\sd\stable-diffusion-webui\repositories\k-diffusion\k_diffusion\external.py", line 112, in forward
        eps = self.get_eps(input * c_in, self.sigma_to_t(sigma), **kwargs)
      File "C:\sd\stable-diffusion-webui\repositories\k-diffusion\k_diffusion\external.py", line 138, in get_eps
        return self.inner_model.apply_model(*args, **kwargs)
      File "C:\sd\stable-diffusion-webui\modules\sd_hijack_utils.py", line 17, in <lambda>
        setattr(resolved_obj, func_path[-1], lambda *args, **kwargs: self(*args, **kwargs))
      File "C:\sd\stable-diffusion-webui\modules\sd_hijack_utils.py", line 28, in __call__
        return self.__orig_func(*args, **kwargs)
      File "C:\sd\stable-diffusion-webui\repositories\stable-diffusion-stability-ai\ldm\models\diffusion\ddpm.py", line 858, in apply_model
        x_recon = self.model(x_noisy, t, **cond)
      File "C:\sd\stable-diffusion-webui\venv\lib\site-packages\torch\nn\modules\module.py", line 1501, in _call_impl
        return forward_call(*args, **kwargs)
      File "C:\sd\stable-diffusion-webui\repositories\stable-diffusion-stability-ai\ldm\models\diffusion\ddpm.py", line 1335, in forward
        out = self.diffusion_model(x, t, context=cc)
      File "C:\sd\stable-diffusion-webui\venv\lib\site-packages\torch\nn\modules\module.py", line 1538, in _call_impl
        result = forward_call(*args, **kwargs)
      File "C:\sd\stable-diffusion-webui\modules\sd_unet.py", line 89, in UNetModel_forward
        return current_unet.forward(x, timesteps, context, *args, **kwargs)
      File "C:\sd\stable-diffusion-webui\extensions\stable-diffusion-webui-tensorrt\scripts\trt.py", line 86, in forward        self.infer({"x": x, "timesteps": timesteps, "context": context})
      File "C:\sd\stable-diffusion-webui\extensions\stable-diffusion-webui-tensorrt\scripts\trt.py", line 72, in infer
        self.buffers[name].copy_(tensor)
    KeyError: 'context'

A 512x512 works fine. Any insight?

Platform:
Windows 10/11
64gb RAM
Python 3.10.6
RTX 3080 10gb/4080 16gb (tried both)
TensorRT-8.6.1.6 GA
Driver Version: 536.25
CUDA Version: 12.2
v1.4.1-172-gff73841c
torch: 2.0.1+cu118

[cd XPHUPZHJ]

same for me
see #84