space running locally on cpu

app.py

@@ -1,41 +1,39 @@
 import gradio as gr
+import numpy as np
 
-from chrislib.general import uninvert, view
+from chrislib.general import uninvert, invert, view, view_scale
 
 from intrinsic.pipeline import load_models, run_pipeline
 
-# load the intrinsic models
-intrinsic_models = load_models('v2')
+intrinsic_models = load_models('v2', device='cpu')
 
 def generate_pipeline(models):
 
-    def pipeline_func(image):
-        return run_pipeline(models, image, device='cuda')
+    def pipeline_func(image, **kwargs):
+        return run_pipeline(models, image, **kwargs)
     
     return pipeline_func
 
 
-run_pipeline = generate_pipeline(intrinsic_models)
+pipeline_func = generate_pipeline(intrinsic_models)
 
 def process_image(image):
     print(image.shape)
+    image = image.astype(np.single) / 255.
 
-    # Process the image with your model pipeline
-    result = run_pipeline(image)
+    result = pipeline_func(image, device='cpu')
     
-    out_keys = ['hr_alb', 'dif_shd', 'pos_res']
-    # Create a list to store output images
-    output_images = []
-    for k in out_keys:
-        output_images.append(result[k])
-
-    return output_images
+    return [view(result['hr_alb']), 1 - invert(result['dif_shd']), view_scale(result['pos_res'])]
 
 interface = gr.Interface(
     fn=process_image,
-    inputs=gr.inputs.Image(type="numpy"),
-    outputs=gr.outputs.Image(type="numpy", label="Output Images", tool="editor"),
+    inputs=gr.components.Image(type="numpy"),
+    outputs=[
+        gr.components.Image(type='numpy', label="Albedo"),
+        gr.components.Image(type='numpy', label="Diffuse Shading"),
+        gr.components.Image(type='numpy', label="Residual"),
+    ],
     live=True
 )
 
-interface.launch()
+interface.launch()