Fixed tokenizer bug and gibberish image generation. Use AutoTokenizer instead of MT5Tokenizer

" Some weights of MT5EncoderModel were not initialized from the model checkpoint at google/mt5-small and are newly initialized: ['encoder.embed_tokens.weight', 'shared.weight']
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference. "

utils.py

@@ -9,16 +9,29 @@ from imwatermark import WatermarkEncoder
 
 from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
 from diffusers.utils.torch_utils import randn_tensor
-from transformers import MT5Tokenizer, MT5EncoderModel
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 from typing import List, Optional, Tuple, Union
 
 # Determine device and torch dtype
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
 
-# Load MT5 tokenizer and encoder (can be replaced with private model + token if needed)
-tokenizer = MT5Tokenizer.from_pretrained("google/mt5-small", use_safetensors=True)
-encoder_model = MT5EncoderModel.from_pretrained("google/mt5-small", use_safetensors=True).to(device=device, dtype=torch_dtype)
+model_name = "google/mt5-small"  # or base / large / etc.
+
+# Load tokenizer:
+tokenizer = AutoTokenizer.from_pretrained(
+    model_name,
+    use_fast=False,       # use slow (SentencePiece) tokenizer to avoid subtle fast/slow differences
+    legacy=True            # enforce legacy behavior (so that “new vs legacy” mismatch warnings are avoided)
+)
+
+# Load model:
+encoder_model = AutoModelForSeq2SeqLM.from_pretrained(
+    model_name,
+    torch_dtype=torch_dtype,     # or whatever dtype you want (float32/float16/bfloat16)
+    device_map="auto"               # or device=device if you want to manually move
+)
+
 encoder_model.eval()
 
 class QPipeline(DiffusionPipeline):