more fixes and cleanup

scripts/demo/streamlit_helpers.py

@@ -35,7 +35,7 @@ def init_st(spec: SamplingSpec, load_ckpt=True, load_filter=True) -> Dict[str, A
             pipeline = SamplingPipeline(
                 model_spec=spec,
                 use_fp16=True,
-                model_loader=CudaModelManager(device="cuda", swap_device="cpu"),
+                device_manager=CudaModelManager(device="cuda", swap_device="cpu"),
             )
         else:
             pipeline = SamplingPipeline(model_spec=spec, use_fp16=False)
@@ -207,7 +207,7 @@ def get_discretization(params: SamplingParams, key=1) -> SamplingParams:
 
 
 def get_sampler(params: SamplingParams, key=1) -> SamplingParams:
-    if params.sampler == Sampler.EULER_EDM or params.sampler == Sampler.HEUN_EDM:
+    if params.sampler in (Sampler.EULER_EDM, Sampler.HEUN_EDM):
         params.s_churn = st.sidebar.number_input(
             f"s_churn #{key}", value=params.s_churn, min_value=0.0
         )
@@ -221,10 +221,7 @@ def get_sampler(params: SamplingParams, key=1) -> SamplingParams:
             f"s_noise #{key}", value=params.s_noise, min_value=0.0
         )
 
-    elif (
-        params.sampler == Sampler.EULER_ANCESTRAL
-        or params.sampler == Sampler.DPMPP2S_ANCESTRAL
-    ):
+    elif params.sampler in (Sampler.EULER_ANCESTRAL, Sampler.DPMPP2S_ANCESTRAL):
         params.s_noise = st.sidebar.number_input(
             "s_noise", value=params.s_noise, min_value=0.0
         )

sgm/inference/api.py

@@ -205,9 +205,9 @@ def __init__(
                 f"Checkpoint {self.ckpt} not found, check model spec or config_path"
             )
 
-        self.model_manager = device_manager
+        self.device_manager = device_manager
         self.model = self._load_model(
-            device_manager=self.model_manager, use_fp16=use_fp16
+            device_manager=self.device_manager, use_fp16=use_fp16
         )
 
     def _load_model(self, device_manager: DeviceModelManager, use_fp16=True):
@@ -229,7 +229,7 @@ def text_to_image(
         samples: int = 1,
         return_latents: bool = False,
         noise_strength: Optional[float] = None,
-        filter: Any = None,
+        filter=None,
     ):
         sampler = get_sampler_config(params)
 
@@ -257,7 +257,7 @@ def text_to_image(
             force_uc_zero_embeddings=["txt"] if not self.specs.is_legacy else [],
             return_latents=return_latents,
             filter=filter,
-            model_manager=self.model_manager,
+            device_manager=self.device_manager,
         )
 
     def image_to_image(
@@ -269,7 +269,7 @@ def image_to_image(
         samples: int = 1,
         return_latents: bool = False,
         noise_strength: Optional[float] = None,
-        filter: Any = None,
+        filter=None,
     ):
         sampler = get_sampler_config(params)
 
@@ -295,7 +295,7 @@ def image_to_image(
             force_uc_zero_embeddings=["txt"] if not self.specs.is_legacy else [],
             return_latents=return_latents,
             filter=filter,
-            device=self.device,
+            device_manager=self.device_manager,
         )
 
     def wrap_discretization(
@@ -364,7 +364,7 @@ def refiner(
             return_latents=return_latents,
             filter=filter,
             add_noise=add_noise,
-            device=self.device,
+            device_manager=self.device_manager,
         )