diff --git a/setup.py b/setup.py
index 374ee85b..c6537d5b 100644
--- a/setup.py
+++ b/setup.py
@@ -13,8 +13,9 @@
 # libavcodec-extra : libavcodec-extra  includes additional codecs for ffmpeg
 
 install_requires = [
-    "transformers[sklearn,sentencepiece,audio,vision]==4.41.1",
-    "orjson",
+    "transformers[sklearn,sentencepiece,audio,vision,sentencepiece]==4.44.0",
+    "huggingface_hub[hf_transfer]==0.24.5",
+    "peft==0.12.0",
     # vision
     "Pillow",
     "librosa",
@@ -26,13 +27,13 @@
     "starlette",
     "uvicorn",
     "pandas",
-    "peft==0.11.1",
+    "orjson",
 ]
 
 extras = {}
 
 extras["st"] = ["sentence_transformers==2.7.0"]
-extras["diffusers"] = ["diffusers==0.26.3", "accelerate==0.27.2"]
+extras["diffusers"] = ["diffusers==0.30.0", "accelerate==0.33.0"]
 extras["torch"] = ["torch==2.2.2", "torchvision", "torchaudio"]
 extras["test"] = [
     "pytest==7.2.1",
@@ -53,7 +54,7 @@
 setup(
     name="huggingface-inference-toolkit",
     version=VERSION,
-    author="HuggingFace",
+    author="Hugging Face",
     description="Hugging Face Inference Toolkit is for serving 🤗 Transformers models in containers.",
     url="",
     package_dir={"": "src"},
diff --git a/src/huggingface_inference_toolkit/diffusers_utils.py b/src/huggingface_inference_toolkit/diffusers_utils.py
index afe96676..54fad5ad 100644
--- a/src/huggingface_inference_toolkit/diffusers_utils.py
+++ b/src/huggingface_inference_toolkit/diffusers_utils.py
@@ -28,7 +28,7 @@ def __init__(
         dtype = torch.float32
         if device == "cuda":
             dtype = torch.bfloat16 if is_torch_bf16_gpu_available() else torch.float16
-        device_map = "auto" if device == "cuda" else None
+        device_map = "balanced" if device == "cuda" else None
 
         self.pipeline = AutoPipelineForText2Image.from_pretrained(
             model_dir, torch_dtype=dtype, device_map=device_map, **kwargs
@@ -42,8 +42,6 @@ def __init__(
             except Exception:
                 pass
 
-        self.pipeline.to(device)
-
     def __call__(
         self,
         prompt,
diff --git a/src/huggingface_inference_toolkit/utils.py b/src/huggingface_inference_toolkit/utils.py
index 89261d71..a5ff7aee 100644
--- a/src/huggingface_inference_toolkit/utils.py
+++ b/src/huggingface_inference_toolkit/utils.py
@@ -8,11 +8,7 @@
 from transformers.file_utils import is_tf_available, is_torch_available
 from transformers.pipelines import Pipeline
 
-from huggingface_inference_toolkit.const import (
-    HF_DEFAULT_PIPELINE_NAME,
-    HF_MODULE_NAME,
-    HF_TRUST_REMOTE_CODE,
-)
+from huggingface_inference_toolkit.const import HF_DEFAULT_PIPELINE_NAME, HF_MODULE_NAME
 from huggingface_inference_toolkit.diffusers_utils import (
     get_diffusers_pipeline,
     is_diffusers_available,
@@ -240,7 +236,7 @@ def get_pipeline(
         "zero-shot-image-classification",
     }:
         kwargs["feature_extractor"] = model_dir
-    elif task in {"image-to-text"}:
+    elif task in {"image-to-text", "text-to-image"}:
         pass
     elif task == "conversational":
         task = "text-generation"