adding model switch

app.py

@@ -9,24 +9,35 @@ from intrinsic.pipeline import load_models, run_pipeline
 
 DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
-intrinsic_models = load_models('v2', device=DEVICE)
+# Instead of loading models at startup, we'll create a cache for models
+model_cache = {}
 
-def generate_pipeline(models):
+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]
 
+def generate_pipeline(models):
     def pipeline_func(image, **kwargs):
         return run_pipeline(models, image, **kwargs)
     
     return pipeline_func
 
-
-pipeline_func = generate_pipeline(intrinsic_models)
-
 @spaces.GPU
-def process_image(image):
+def process_image(image, model_version):
+    # Check if image is provided
+    if image is None:
+        return [None, None, None]
+        
+    print(f"Processing with model version: {model_version}")
     print(image.shape)
     image = image.astype(np.single) / 255.
 
-    result = pipeline_func(image, device=DEVICE)
+    # Get or load the selected model
+    models = get_model(model_version)
+    pipeline_func = generate_pipeline(models)
+    
+    result = pipeline_func(image, device=DEVICE, resize_conf=1024)
     
     return [view(result['hr_alb']), 1 - invert(result['dif_shd']), view_scale(result['pos_res'])]
 
@@ -61,6 +72,16 @@ with gr.Blocks(
         .md_feedback li {
             margin-bottom: 0px !important;
         }
+        .image-gallery {
+            display: flex;
+            flex-wrap: wrap;
+            gap: 10px;
+            justify-content: center;
+        }
+        .image-gallery > * {
+            flex: 1;
+            min-width: 200px;
+        }
         """,
 ) as demo:
     gr.Markdown(
@@ -75,14 +96,34 @@ with gr.Blocks(
             </a>
         """
     )
-    with gr.Row():
+    
+    # More compact model version selector using dropdown
+    model_version = gr.Dropdown(
+        choices=["v2", "v2.1"], 
+        value="v2", 
+        label="Model Version",
+        info="Select which model weights to use",
+        scale=2
+    )
+    
+    # Gallery-style layout for all images
+    with gr.Row(elem_classes="image-gallery"):
         input_img = gr.Image(label="Input Image")
-
-    with gr.Row():
         alb_img = gr.Image(label="Albedo")
         shd_img = gr.Image(label="Diffuse Shading")
         dif_img = gr.Image(label="Diffuse Image")
 
-    input_img.change(process_image, inputs=input_img, outputs=[alb_img, shd_img, dif_img])
+    # Update to pass model_version to process_image
+    input_img.change(
+        process_image, 
+        inputs=[input_img, model_version], 
+        outputs=[alb_img, shd_img, dif_img]
+    )
+    # Add event handler for when model_version changes
+    model_version.change(
+        process_image,
+        inputs=[input_img, model_version],
+        outputs=[alb_img, shd_img, dif_img]
+    )
 
 demo.launch(show_error=True)