Possibility to override dtype trhough env var

Signed-off-by: Raphael Glon <[email protected]>

Author: oOraph

docker_images/diffusers/app/pipelines/image_to_image.py

@@ -45,7 +45,10 @@ def __init__(self, model_id: str):
             if model_id.startswith("hf-internal-testing/")
             else {}
         )
-        if torch.cuda.is_available():
+        env_dtype = os.getenv("TORCH_DTYPE")
+        if env_dtype:
+            kwargs["torch_dtype"] = getattr(torch, env_dtype)
+        elif torch.cuda.is_available():
             kwargs["torch_dtype"] = torch.float16
             if model_id == "stabilityai/stable-diffusion-xl-refiner-1.0":
                 kwargs["variant"] = "fp16"

docker_images/diffusers/app/pipelines/text_to_image.py

@@ -40,7 +40,10 @@ def __init__(self, model_id: str):
             if model_id.startswith("hf-internal-testing/")
             else {}
         )
-        if torch.cuda.is_available():
+        env_dtype = os.getenv("TORCH_DTYPE")
+        if env_dtype:
+            kwargs["torch_dtype"] = getattr(torch, env_dtype)
+        elif torch.cuda.is_available():
             kwargs["torch_dtype"] = torch.float16
 
         has_model_index = any(