loading models at start up

app.py

@@ -12,10 +12,8 @@ DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 # Instead of loading models at startup, we'll create a cache for models
 model_cache = {}
 
-def get_model(model_version):
-    if model_version not in model_cache:
-        model_cache[model_version] = load_models(model_version, device=DEVICE)
-    return model_cache[model_version]
+model_cache['v2'] = load_models('v2', device=DEVICE)
+model_cache['v2.1'] = load_models('v2.1', device=DEVICE)
 
 def generate_pipeline(models):
     def pipeline_func(image, **kwargs):
@@ -34,7 +32,7 @@ def process_image(image, model_version):
     image = image.astype(np.single) / 255.
 
     # Get or load the selected model
-    models = get_model(model_version)
+    models = model_cache[model_version]
     pipeline_func = generate_pipeline(models)
     
     result = pipeline_func(image, device=DEVICE, resize_conf=1024)
@@ -87,13 +85,6 @@ with gr.Blocks(
     gr.Markdown(
             """
             # Colorful Diffuse Intrinsic Image Decomposition in the Wild
-            <p align="center">
-            <a title="Website" href="https://yaksoy.github.io/ColorfulShading/" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
-                <img src="https://www.obukhov.ai/img/badges/badge-website.svg">
-            </a>
-            <a title="Github" href="https://github.com/compphoto/Intrinsic" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
-                <img src="https://img.shields.io/github/stars/compphoto/Intrinsic?label=GitHub%20%E2%98%85&logo=github&color=C8C" alt="badge-github-stars">
-            </a>
         """
     )
     
@@ -108,8 +99,17 @@ with gr.Blocks(
         )
         
         gr.Markdown("""
-        The model may take a few seconds to load the first time you use it.
-        Subsequent decompositions should be faster after the model is loaded.
+            <p align="center">
+            <a title="Website" href="https://yaksoy.github.io/ColorfulShading/" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
+                <img src="https://www.obukhov.ai/img/badges/badge-website.svg">
+            </a>
+            <a title="Github" href="https://github.com/compphoto/Intrinsic" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
+                <img src="https://img.shields.io/github/stars/compphoto/Intrinsic?label=GitHub%20%E2%98%85&logo=github&color=C8C" alt="badge-github-stars">
+            </a>
+
+            **V2**: Original weights from the paper.
+            
+            **V2.1**: More albedo detail and improved diffuse shading estimation.
         """)
     
     # Gallery-style layout for all images