change to image to imagw

app_tensorrt.py

@@ -50,39 +50,51 @@ compiledModel = torch.compile(
 
 base_pipe.unet = compiledModel
 
-init_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/img2img.png")
-generated_image = base_pipe(
-    image=init_image,
-    prompt="A white cat",
-    num_inference_steps=5,
-).images[0]
+import torch._dynamo
+torch._dynamo.config.suppress_errors = True
 
-generated_image.save("/tmp/gradio/generated_image.png")
+try:
+    init_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/img2img.png")
+    generated_image = base_pipe(
+        image=init_image,
+        prompt="A white cat",
+        num_inference_steps=5,
+    ).images[0]
+
+    generated_image.save("/tmp/gradio/generated_image.png")
+except Exception as e:
+    print(f"Error: {e}")
 
 
 def create_demo() -> gr.Blocks:
 
     @spaces.GPU(duration=30)
-    def text_to_image(
+    def image_to_image(
+        image: gr.Image,
         prompt:str,
         steps:int,
     ):
-        import torch_tensorrt
-        print('Compiling model...')
-        compiledModel = torch.compile(
-            base_pipe.unet,
-            backend=backend,
-            options={
-                "truncate_long_and_double": True,
-                "enabled_precisions": {torch.float32, torch.float16},
-            },
-            dynamic=False,
-        )
-        print('Model compiled!')
-
-        print('Saving compiled model...')
-        torch_tensorrt.save(compiledModel, "compiled_pipe.ep")
-        print('Compiled model saved!')
+        run_task_time = 0
+        time_cost_str = ''
+        run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)
+        generated_image = base_pipe(
+            image=image,
+            prompt=prompt,
+            num_inference_steps=steps,
+        ).images[0]
+        run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)
+        return generated_image
+    
+    def get_time_cost(run_task_time, time_cost_str):
+        now_time = int(time.time()*1000)
+        if run_task_time == 0:
+            time_cost_str = 'start'
+        else:
+            if time_cost_str != '': 
+                time_cost_str += f'-->'
+            time_cost_str += f'{now_time - run_task_time}'
+        run_task_time = now_time
+        return run_task_time, time_cost_str
 
     with gr.Blocks() as demo:
         with gr.Row():
@@ -94,15 +106,15 @@ def create_demo() -> gr.Blocks:
                 
         with gr.Row():
             with gr.Column():
-                generated_image = gr.Image(label="Generated Image", type="pil", interactive=False)
+                input_image = gr.Image(label="Input Image", type="pil", interactive=True)
             with gr.Column():
+                generated_image = gr.Image(label="Generated Image", type="pil", interactive=False)
                 time_cost = gr.Textbox(label="Time Cost", lines=1, interactive=False)
         
         g_btn.click(
             fn=text_to_image,
-            inputs=[prompt, steps],
-            # outputs=[generated_image, time_cost],
-            outputs=[],
+            inputs=[input_image, prompt, steps],
+            outputs=[generated_image, time_cost],
         )
 
     return demo