refactor: apply small code improvements (#113)

This commit makes some very small code improvements that make applying a mock patch easier.
livepeer · Jun 26, 2024 · e5f4bf7 · e5f4bf7
1 parent 8b1b455
commit e5f4bf7
Show file tree

Hide file tree

Showing 4 changed files with 4 additions and 5 deletions.
diff --git a/runner/app/pipelines/image_to_image.py b/runner/app/pipelines/image_to_image.py
@@ -45,6 +45,7 @@ def list(cls):
 
 class ImageToImagePipeline(Pipeline):
     def __init__(self, model_id: str):
+        self.model_id = model_id
         kwargs = {"cache_dir": get_model_dir()}
 
         torch_device = get_torch_device()
@@ -69,8 +70,6 @@ def __init__(self, model_id: str):
             kwargs["torch_dtype"] = torch.float16
             kwargs["variant"] = "fp16"
 
-        self.model_id = model_id
-
         # Special case SDXL-Lightning because the unet for SDXL needs to be swapped
         if ModelName.SDXL_LIGHTNING.value in model_id:
             base = "stabilityai/stable-diffusion-xl-base-1.0"

diff --git a/runner/app/pipelines/image_to_video.py b/runner/app/pipelines/image_to_video.py
@@ -21,6 +21,7 @@
 
 class ImageToVideoPipeline(Pipeline):
     def __init__(self, model_id: str):
+        self.model_id = model_id
         kwargs = {"cache_dir": get_model_dir()}
 
         torch_device = get_torch_device()
@@ -39,7 +40,6 @@ def __init__(self, model_id: str):
             kwargs["torch_dtype"] = torch.float16
             kwargs["variant"] = "fp16"
 
-        self.model_id = model_id
         self.ldm = StableVideoDiffusionPipeline.from_pretrained(model_id, **kwargs)
         self.ldm.to(get_torch_device())
 

diff --git a/runner/app/pipelines/text_to_image.py b/runner/app/pipelines/text_to_image.py
@@ -23,6 +23,7 @@
 
 class TextToImagePipeline(Pipeline):
     def __init__(self, model_id: str):
+        self.model_id = model_id
         kwargs = {"cache_dir": get_model_dir()}
 
         torch_device = get_torch_device()
@@ -51,7 +52,6 @@ def __init__(self, model_id: str):
             logger.info("TextToImagePipeline using bfloat16 precision for %s", model_id)
             kwargs["torch_dtype"] = torch.bfloat16
 
-        self.model_id = model_id
 
         # Special case SDXL-Lightning because the unet for SDXL needs to be swapped
         if SDXL_LIGHTNING_MODEL_ID in model_id:

diff --git a/runner/app/pipelines/upscale.py b/runner/app/pipelines/upscale.py
@@ -24,6 +24,7 @@
 
 class UpscalePipeline(Pipeline):
     def __init__(self, model_id: str):
+        self.model_id = model_id
         kwargs = {"cache_dir": get_model_dir()}
 
         torch_device = get_torch_device()
@@ -42,7 +43,6 @@ def __init__(self, model_id: str):
             kwargs["torch_dtype"] = torch.float16
             kwargs["variant"] = "fp16"
 
-        self.model_id = model_id
         self.ldm = StableDiffusionUpscalePipeline.from_pretrained(
                 model_id, **kwargs
             ).to(torch_device)