test

app.py

@@ -1,4 +1,4 @@
-# Copyright 2024 Anton Obukhov, ETH Zurich. All rights reserved.
+# Copyright 2024 Anton Obukhov and Kevin Qu, ETH Zurich. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -57,8 +57,8 @@ loaded_pipelines = {}  # Cache to store loaded pipelines
 def process_with_loaded_pipeline(image_path, denoise_steps, ensemble_size, processing_res, model_type):
 
     # Load and cache the pipeline based on the model type.
-    if model_type not in loaded_pipelines:
-        auth_token = os.environ.get("KEV_TOKEN")
+    if model_type not in loaded_pipelines.keys():
+        auth_token = os.environ.get("KEV_DEV")
         if model_type == "appearance":
             loaded_pipelines[model_type] = MarigoldIIDAppearancePipeline.from_pretrained(
                 "prs-eth/marigold-iid-appearance-v1-1", token=auth_token
@@ -71,6 +71,12 @@ def process_with_loaded_pipeline(image_path, denoise_steps, ensemble_size, proce
         # Move the pipeline to GPU if available
         device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         loaded_pipelines[model_type] = loaded_pipelines[model_type].to(device)
+        try:
+            import xformers
+
+            loaded_pipelines[model_type].enable_xformers_memory_efficient_attention()
+        except:
+            pass  # run without xformers
     
     pipe = loaded_pipelines[model_type]
 
@@ -511,6 +517,9 @@ def run_demo_server(hf_writer=None):
                 None,
                 None,
                 None,
+                None,
+                None,
+                default_model_type,
                 default_image_ensemble_size,
                 default_image_denoise_steps,
                 default_image_processing_res,
@@ -522,6 +531,7 @@ def run_demo_server(hf_writer=None):
                 image_output_slider2,
                 image_output_slider3,
                 image_output_files,
+                model_type,
                 image_ensemble_size,
                 image_denoise_steps,
                 image_processing_res,
@@ -602,25 +612,12 @@ def run_demo_server(hf_writer=None):
 
 
 def main():
-    CHECKPOINT = "prs-eth/marigold-iid-appearance-v1-1"
     CROWD_DATA = "crowddata-marigold-iid-appearance-v1-1-space-v1-1"
 
     os.system("pip freeze")
 
     if "HF_TOKEN_LOGIN" in os.environ:
         login(token=os.environ["HF_TOKEN_LOGIN"])
-    
-    auth_token = os.environ.get("KEV_TOKEN") 
-    pipe = MarigoldIIDAppearancePipeline.from_pretrained(CHECKPOINT,token=auth_token)
-    try:
-        import xformers
-
-        pipe.enable_xformers_memory_efficient_attention()
-    except:
-        pass  # run without xformers
-
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    pipe = pipe.to(device)
 
     hf_writer = None
     if "HF_TOKEN_LOGIN_WRITE_CROWD" in os.environ: