Improve app

app.py

@@ -19,23 +19,54 @@ class Examples(gr.helpers.Examples):
             self.cached_file = Path(self.cached_folder) / "log.csv"
         self.create()
 
-def load_predictor():
-    """Load model predictor using torch.hub"""
-    predictor = torch.hub.load("Stable-X/StableNormal", "StableNormal_turbo", 
-                               trust_repo=True, yoso_version='yoso-normal-v1-8-1')
-    return predictor
+# Global variable to store loaded predictors
+predictors = {}
+
+# Available model versions
+MODEL_VERSIONS = {
+    "v0.3": "yoso-normal-v0-3",
+    "v1.0": "yoso-normal-v1-0", 
+    "v1.5": "yoso-normal-v1-5",
+    "v1.8.1": "yoso-normal-v1-8-1"
+}
+
+def load_predictor(version: str = "v1.8.1"):
+    """Load model predictor using torch.hub with specified version"""
+    if version not in predictors:
+        yoso_version = MODEL_VERSIONS[version]
+        print(f"Loading StableNormal with {yoso_version}...")
+        predictor = torch.hub.load("Stable-X/StableNormal", "StableNormal_turbo", 
+                                   trust_repo=True, yoso_version=yoso_version)
+        predictors[version] = predictor
+        print(f"Successfully loaded {version}")
+    return predictors[version]
+
+def precache_all_predictors():
+    """Precache all model predictors at startup"""
+    print("Precaching all StableNormal predictors...")
+    for version in MODEL_VERSIONS.keys():
+        print(f"Precaching {version}...")
+        try:
+            load_predictor(version)
+            print(f"✓ Successfully precached {version}")
+        except Exception as e:
+            print(f"✗ Failed to precache {version}: {e}")
+    print("Finished precaching all predictors.")
 
 def process_image(
-    predictor,
     path_input: str,
+    version: str = "v1.8.1",
     data_type: str = "object"
 ) -> tuple:
-    """Process single image"""
+    """Process single image with specified model version"""
     if path_input is None:
         raise gr.Error("Please upload an image or select one from the gallery.")
+    
+    # Load the predictor for the specified version
+    predictor = load_predictor(version)
         
     name_base = os.path.splitext(os.path.basename(path_input))[0]
-    out_path = os.path.join(tempfile.mkdtemp(), f"{name_base}_normal.png")
+    out_path = os.path.join(tempfile.mkdtemp(), f"{name_base}_normal_{version.replace('.', '_')}.png")
 
     # Load and process image
     input_image = Image.open(path_input)
@@ -45,16 +76,15 @@ def process_image(
     yield [input_image, out_path]
 
 def create_demo():
-    # Load model
-    predictor = load_predictor()
+    # Precache all predictors before creating the demo
+    precache_all_predictors()
     
-    # Create processing functions for each data type
-    process_object = spaces.GPU(functools.partial(process_image, predictor, data_type="object"))
+    # Create processing function
+    process_object = spaces.GPU(process_image)
 
     # Define markdown content
     HEADER_MD = """
     # 🎪 StableNormal Turbo
-
     <p align="center">
     <a title="Website" href="https://stable-x.github.io/StableNormal/" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
         <img src="https://www.obukhov.ai/img/badges/badge-website.svg">
@@ -69,6 +99,8 @@ def create_demo():
         <img src="https://www.obukhov.ai/img/badges/badge-social.svg" alt="social">
     </a>
     </p>
+    
+    Select between different YOSO Normal model versions. Each version may have different performance characteristics and quality trade-offs.
     """
 
     # Create interface
@@ -92,9 +124,19 @@ def create_demo():
                 with gr.Row():
                     with gr.Column():
                         object_input = gr.Image(label="Input Object Image", type="filepath")
+                        
+                        # Model version selector
+                        version_dropdown = gr.Dropdown(
+                            choices=list(MODEL_VERSIONS.keys()),
+                            value="v1.8.1",
+                            label="Model Version",
+                            info="Select YOSO Normal model version"
+                        )
+                        
                         with gr.Row():
                             object_submit_btn = gr.Button("Compute Normal", variant="primary")
                             object_reset_btn = gr.Button("Reset")
+                            
                     with gr.Column():
                         object_output_slider = ImageSlider(
                             label="Normal outputs",
@@ -106,36 +148,49 @@ def create_demo():
                             position=0.25,
                         )
 
-                Examples(
-                    fn=process_object,
-                    examples=sorted([
-                        os.path.join("files", "object", name)
-                        for name in os.listdir(os.path.join("files", "object"))
-                        if os.path.exists(os.path.join("files", "object"))
-                    ]),
-                    inputs=[object_input],
-                    outputs=[object_output_slider],
-                    cache_examples=False,
-                    directory_name="examples_object",
-                    examples_per_page=50,
-                )
+                # Model version info
+                with gr.Row():
+                    gr.Markdown("""
+                    **Model Version Information:**
+                    - **v0.3**: Camera Ready Version
+                    - **v1.0**: Improve stability, but poor sharpness
+                    - **v1.5**: Enhanced performance and accuracy  
+                    - **v1.8.1**: Latest version with best sharpness (default)
+                    
+                    *All models are precached and ready for instant switching.*
+                    """)
+
+                # Examples section
+                if os.path.exists(os.path.join("files", "object")):
+                    Examples(
+                        fn=lambda img, ver: process_object(img, ver),
+                        examples=sorted([
+                            [os.path.join("files", "object", name), "v1.8.1"]
+                            for name in os.listdir(os.path.join("files", "object"))
+                        ]),
+                        inputs=[object_input, version_dropdown],
+                        outputs=[object_output_slider],
+                        cache_examples=False,
+                        directory_name="examples_object",
+                        examples_per_page=50,
+                    )
 
         # Event Handlers for Object Tab
         object_submit_btn.click(
-            fn=lambda x, _: None if x else gr.Error("Please upload an image"),
-            inputs=object_input,
+            fn=lambda x, v: None if x else gr.Error("Please upload an image"),
+            inputs=[object_input, version_dropdown],
             outputs=None,
             queue=False,
         ).success(
             fn=process_object,
-            inputs=object_input,
+            inputs=[object_input, version_dropdown],
             outputs=[object_output_slider],
         )
 
         object_reset_btn.click(
-            fn=lambda: (None, DEFAULT_SHARPNESS, None),
+            fn=lambda: (None, "v1.8.1", None),
             inputs=[],
-            outputs=[object_input, object_output_slider],
+            outputs=[object_input, version_dropdown, object_output_slider],
             queue=False,
         )
 

requirements.txt

@@ -30,9 +30,9 @@ filelock==3.14.0
 fonttools==4.53.0
 frozenlist==1.4.1
 fsspec==2024.3.1
-gradio==4.32.2
-gradio_client==0.17.0
-gradio_imageslider==0.0.20
+gradio==4.44.1
+gradio_client
+gradio_imageslider
 h11==0.14.0
 httpcore==1.0.5
 httptools==0.6.1