Rebased to main after refit acceleration is merged.

Author: cehongwang

examples/dynamo/mutable_torchtrt_module_example.py

@@ -21,8 +21,8 @@
 import torch_tensorrt as torch_trt
 import torchvision.models as models
 
-np.random.seed(0)
-torch.manual_seed(0)
+np.random.seed(5)
+torch.manual_seed(5)
 inputs = [torch.rand((1, 3, 224, 224)).to("cuda")]
 
 # %%
@@ -76,7 +76,7 @@
 from diffusers import DiffusionPipeline
 
 with torch.no_grad():
-    kwargs = {
+    settings = {
         "use_python_runtime": True,
         "enabled_precisions": {torch.float16},
         "debug": True,
@@ -95,7 +95,7 @@
     pipe.to(device)
 
     # The only extra line you need
-    pipe.unet = torch_trt.MutableTorchTensorRTModule(pipe.unet, **kwargs)
+    pipe.unet = torch_trt.MutableTorchTensorRTModule(pipe.unet, **settings)
 
     image = pipe(prompt, negative_prompt=negative, num_inference_steps=30).images[0]
     image.save("./without_LoRA_mutable.jpg")

py/torch_tensorrt/dynamo/runtime/_MutableTorchTensorRTModule.py

@@ -267,7 +267,9 @@ def refit_gm(self) -> None:
                     self.original_model.state_dict()
                 )
             )
-        self.gm = refit_module_weights(self.gm, self.exp_program)
+        self.gm = refit_module_weights(
+            self.gm, self.exp_program, use_weight_map_cache=True, in_place=True
+        )
 
         self.original_model.cpu()
         torch.cuda.empty_cache()