multiple values for keyword argument 'opt_batch_size'

[olegchomp]

trafficstars

Got this error always when trying to call accelerate_with_tensorrt.

import torch
from diffusers import AutoencoderTiny, StableDiffusionPipeline
from diffusers.utils import load_image

from streamdiffusion import StreamDiffusion
from streamdiffusion.image_utils import postprocess_image
from streamdiffusion.acceleration.tensorrt import accelerate_with_tensorrt

# You can load any models using diffuser's StableDiffusionPipeline
pipe = StableDiffusionPipeline.from_pretrained("KBlueLeaf/kohaku-v2.1").to(
    device=torch.device("cuda"),
    dtype=torch.float16,
)

# Wrap the pipeline in StreamDiffusion
stream = StreamDiffusion(
    pipe,
    t_index_list=[18,26,35,45],
    torch_dtype=torch.float16,
)

# If the loaded model is not LCM, merge LCM
stream.load_lcm_lora()
stream.fuse_lora()
# Use Tiny VAE for further acceleration
stream.vae = AutoencoderTiny.from_pretrained("madebyollin/taesd").to(device=pipe.device, dtype=pipe.dtype)
# Enable acceleration
#pipe.enable_xformers_memory_efficient_attention()
stream = accelerate_with_tensorrt(
    stream, "engines", max_batch_size=2,
)


prompt = "space banana"
# Prepare the stream
stream.prepare(prompt)

# Prepare image
init_image = load_image("over1.jpeg").resize((512, 512))

# Warmup >= len(t_index_list) x frame_buffer_size
for _ in range(4):
    stream(init_image)

# Run the stream infinitely
while True:
    x_output = stream(init_image)
    postprocess_image(x_output, output_type="pil")[0].show()
    input_response = input("Press Enter to continue or type 'stop' to exit: ")
    if input_response == "stop":
        break

File "streamdiffusion\lib\site-packages\streamdiffusion\acceleration\tensorrt\__init__.py", line 137, in accelerate_with_tensorrt
    compile_unet(
TypeError: streamdiffusion.acceleration.tensorrt.compile_unet() got multiple values for keyword argument 'opt_batch_size'