test generated_image

app_tensorrt.py

@@ -45,16 +45,15 @@ compiledModel = torch.compile(
     },
     dynamic=False,
 )
-print('Model compiled!')
 
-print('Saving compiled model...')
-torch_tensorrt.save(compiledModel, "compiled_pipe.ep")
-print('Compiled model saved!')
+base_pipe.unet = compiledModel
 
-print('Loading compiled model...')
-loadedModel = torch_tensorrt.load("compiled_pipe.ep").module()
-print('Compiled model loaded!')
+generated_image = base_pipe(
+    prompt="A beautiful sunset over the city",
+    num_inference_steps=5,
+).images[0]
 
+generated_image.save("/tmp/gradio/generated_image.png")
 
 
 def create_demo() -> gr.Blocks: