Update app.py

app.py

@@ -71,8 +71,6 @@ def download_models():
     for model, (url, folder, filename) in models.items():
         download_file(url, folder, filename)
 
-
-
 def timer_func(func):
     def wrapper(*args, **kwargs):
         start_time = time.time()
@@ -259,6 +257,7 @@ def get_segmentation_pipeline(
 
 @torch.inference_mode()
 @spaces.GPU
+@timer_func
 def segment_image(
         image: Image,
         image_processor: AutoImageProcessor,
@@ -304,6 +303,7 @@ def get_depth_pipeline():
 
 @torch.inference_mode()
 @spaces.GPU
+@timer_func
 def get_depth_image(
         image: Image,
         feature_extractor: AutoImageProcessor,
@@ -372,6 +372,7 @@ class ControlNetDepthDesignModelMulti:
         self.additional_quality_suffix = "interior design, 4K, high resolution, photorealistic"
         
     @spaces.GPU
+    @timer_func
     def generate_design(self, empty_room_image: Image, prompt: str, guidance_scale: int = 10, num_steps: int = 50, strength: float =0.9, img_size: int = 640) -> Image:
         """
         Given an image of an empty room and a prompt
@@ -552,7 +553,7 @@ def main():
     with gr.Blocks() as demo:
         gr.Markdown(title)
         gr.Markdown(description)   
-        create_demo()
+        create_demo(model)
 
 
     demo.queue().launch(share=False)