add debug log

app.py

@@ -10,6 +10,8 @@ from PIL import Image
 from diffusers import FluxInpaintPipeline
 from huggingface_hub import login
 import os
+import time
+
 
 MARKDOWN = """
 # FLUX.1 Inpainting with lora
@@ -20,9 +22,35 @@ IMAGE_SIZE = 1024
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 HF_TOKEN = os.environ.get("HF_TOKEN")
 trigger_word = "a_photo_of_TOK"
+lora_path = "jiuface/boy-001"
+lora_weights = "flux_train_replicate.safetensors"
+lora_scale = 0.9
 
 login(token=HF_TOKEN)
 
+
+
+class calculateDuration:
+    def __init__(self, activity_name=""):
+        self.activity_name = activity_name
+
+    def __enter__(self):
+        self.start_time = time.time()
+        self.start_time_formatted = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(self.start_time))
+        print(f"Start time: {self.start_time_formatted}")
+        return self
+    
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.end_time = time.time()
+        self.elapsed_time = self.end_time - self.start_time
+        self.end_time_formatted = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(self.end_time))
+        print(f"End time: {self.start_time_formatted}")
+        if self.activity_name:
+            print(f"Elapsed time for {self.activity_name}: {self.elapsed_time:.6f} seconds")
+        else:
+            print(f"Elapsed time: {self.elapsed_time:.6f} seconds")
+
+
 def remove_background(image: Image.Image, threshold: int = 50) -> Image.Image:
     image = image.convert("RGBA")
     data = image.getdata()
@@ -37,34 +65,6 @@ def remove_background(image: Image.Image, threshold: int = 50) -> Image.Image:
     image.putdata(new_data)
     return image
 
-
-EXAMPLES = [
-    [
-        {
-            "background": Image.open(requests.get("https://media.roboflow.com/spaces/doge-2-image.png", stream=True).raw),
-            "layers": [remove_background(Image.open(requests.get("https://media.roboflow.com/spaces/doge-2-mask-2.png", stream=True).raw))],
-            "composite": Image.open(requests.get("https://media.roboflow.com/spaces/doge-2-composite-2.png", stream=True).raw),
-        },
-        "little lion",
-        42,
-        False,
-        0.85,
-        30
-    ],
-    [
-        {
-            "background": Image.open(requests.get("https://media.roboflow.com/spaces/doge-2-image.png", stream=True).raw),
-            "layers": [remove_background(Image.open(requests.get("https://media.roboflow.com/spaces/doge-2-mask-3.png", stream=True).raw))],
-            "composite": Image.open(requests.get("https://media.roboflow.com/spaces/doge-2-composite-3.png", stream=True).raw),
-        },
-        "tribal tattoos",
-        42,
-        False,
-        0.85,
-        30
-    ]
-]
-
 pipe = FluxInpaintPipeline.from_pretrained(
     "black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16).to(DEVICE)
 
@@ -119,27 +119,30 @@ def process(
         gr.Info("Please draw a mask on the image.")
         return None, None
 
-    width, height = resize_image_dimensions(original_resolution_wh=image.size)
-    resized_image = image.resize((width, height), Image.LANCZOS)
-    resized_mask = mask.resize((width, height), Image.LANCZOS)
-
-    pipe.load_lora_weights("jiuface/boy-001", weight_name="flux_train_replicate.safetensors")
-    lora_scale = 0.8
+    with calculateDuration("resize image"):
+        width, height = resize_image_dimensions(original_resolution_wh=image.size)
+        resized_image = image.resize((width, height), Image.LANCZOS)
+        resized_mask = mask.resize((width, height), Image.LANCZOS)
+    with calculateDuration("load lora"):
+        pipe.load_lora_weights(lora_path, weight_name=lora_weights)
+    
     if randomize_seed_checkbox:
         seed_slicer = random.randint(0, MAX_SEED)
     generator = torch.Generator().manual_seed(seed_slicer)
-    result = pipe(
-        prompt=f"{input_text} {trigger_word}",
-        image=resized_image,
-        mask_image=resized_mask,
-        width=width,
-        height=height,
-        strength=strength_slider,
-        generator=generator,
-        num_inference_steps=num_inference_steps_slider,
-        joint_attention_kwargs={"scale": lora_scale},
-    ).images[0]
-    print('INFERENCE DONE')
+
+    with calculateDuration("run pipe"):
+        result = pipe(
+            prompt=f"{input_text} {trigger_word}",
+            image=resized_image,
+            mask_image=resized_mask,
+            width=width,
+            height=height,
+            strength=strength_slider,
+            generator=generator,
+            num_inference_steps=num_inference_steps_slider,
+            joint_attention_kwargs={"scale": lora_scale},
+        ).images[0]
+    
     return result, resized_mask
 
 
@@ -205,25 +208,6 @@ with gr.Blocks() as demo:
             with gr.Accordion("Debug", open=False):
                 output_mask_component = gr.Image(
                     type='pil', image_mode='RGB', label='Input mask', format="png")
-    with gr.Row():
-        gr.Examples(
-            fn=process,
-            examples=EXAMPLES,
-            inputs=[
-                input_image_editor_component,
-                input_text_component,
-                seed_slicer_component,
-                randomize_seed_checkbox_component,
-                strength_slider_component,
-                num_inference_steps_slider_component
-            ],
-            outputs=[
-                output_image_component,
-                output_mask_component
-            ],
-            run_on_click=True,
-            cache_examples=False
-        )
 
     submit_button_component.click(
         fn=process,