replicate/autocompile

replicate/autocompile

Description: automatically infer dynamic shapes for Torch-Tensor compilation

Language: Python

License: MIT

Stars: 4

Forks: 2

Open issues: 0

Created: 2024-10-07T22:04:01Z

Pushed: 2024-11-04T19:58:44Z

Default branch: main

Fork: no

Archived: yes

README:

autocompile

Use tracing to automatically infer static and dynamic shapes as well as which modules need to be compiled in nested pipelies objects, then compile with Torch-TensorRT.

Example usage:

In [1]: from diffusers import StableDiffusionPipeline
...: import tracer

In [2]: pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4")
Loading pipeline components...: 100%|████████████████████████████████████████████████████| 7/7 [00:00Outputs for flux

{ "t5": [], "clip": [], "flux": [ { "min_shape": (1, 3808, 64), "max_shape": (1, 4096, 64), "dtype": torch.bfloat16, }, { "min_shape": (1, 3808, 3), "max_shape": (1, 4096, 3), "dtype": torch.float32, }, {"shape": (1, 256, 4096), "dtype": torch.bfloat16}, {"shape": (1, 256, 3), "dtype": torch.float32}, {"shape": (1,), "dtype": torch.bfloat16}, {"shape": (1, 768), "dtype": torch.bfloat16}, {"shape": (1,), "dtype": torch.bfloat16}, ], "ae.decoder": [ { "min_shape": (1, 16, 112, 112), "max_shape": (1, 16, 136, 136), "dtype": torch.float32, } ], }

{ "t5": {"forward": [{"text": ["a cool dog"]}, {"text": ["a cool dog"]}]}, "t5.hf_module": { "forward": [ { "input_ids": ShapeInfo(shape=torch.Size([1, 256]), dtype=torch.int64), "output_hidden_states": False, }, { "input_ids": ShapeInfo(shape=torch.Size([1, 256]), dtype=torch.int64), "output_hidden_states": False, }, ] }, ... }