Update app.py

app.py

@@ -7,6 +7,7 @@ import spaces
 from diffusers import DiffusionPipeline
 import copy
 import random
+import time
 
 # Load LoRAs from JSON file
 with open('loras.json', 'r') as f:
@@ -19,6 +20,23 @@ pipe.to("cuda")
 
 MAX_SEED = 2**32-1
 
+class calculateDuration:
+    def __init__(self, activity_name=""):
+        self.activity_name = activity_name
+
+    def __enter__(self):
+        self.start_time = time.time()
+        return self
+    
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.end_time = time.time()
+        self.elapsed_time = self.end_time - self.start_time
+        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 update_selection(evt: gr.SelectData):
     selected_lora = loras[evt.index]
     new_placeholder = f"Type a prompt for {selected_lora['title']}"
@@ -40,30 +58,34 @@ def run_lora(prompt, cfg_scale, steps, selected_index, randomize_seed, seed, wid
     trigger_word = selected_lora["trigger_word"]
 
     # Load LoRA weights
-    if "weights" in selected_lora:
-        pipe.load_lora_weights(lora_path, weight_name=selected_lora["weights"])
-    else:
-        pipe.load_lora_weights(lora_path)
+    with calculateDuration("Loading LoRA weights"):
+        if "weights" in selected_lora:
+            pipe.load_lora_weights(lora_path, weight_name=selected_lora["weights"])
+        else:
+            pipe.load_lora_weights(lora_path)
         
     # Set random seed for reproducibility
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-    generator = torch.Generator(device="cuda").manual_seed(seed)
-
-    # Generate image
-    image = pipe(
-        prompt=f"{prompt} {trigger_word}",
-        num_inference_steps=steps,
-        guidance_scale=cfg_scale,
-        width=width,
-        height=height,
-        generator=generator,
-        joint_attention_kwargs={"scale": lora_scale},
-    ).images[0]
+    with calculateDuration("Randomizing seed"):
+        if randomize_seed:
+            seed = random.randint(0, MAX_SEED)
+        generator = torch.Generator(device="cuda").manual_seed(seed)
+
+    with calculateDuration("Generating image"):
+        # Generate image
+        image = pipe(
+            prompt=f"{prompt} {trigger_word}",
+            num_inference_steps=steps,
+            guidance_scale=cfg_scale,
+            width=width,
+            height=height,
+            generator=generator,
+            joint_attention_kwargs={"scale": lora_scale},
+        ).images[0]
 
     yield image, seed
     
-    pipe.unload_lora_weights()
+    with calculateDuration("Unloading weights"):
+        pipe.unload_lora_weights()
     
 
 css = '''