Update dd.py

dd.py

@@ -22,8 +22,8 @@ def format_text_input(prompts: List[str], system_message: str = None):
 
 
 def get_mistral_3_small_prompt_embeds(
-    text_encoder: Mistral3ForConditionalGeneration,
-    tokenizer: AutoProcessor,
+    text_encoder: AutoModelForCausalLM,
+    tokenizer: AutoTokenizer,
     prompt: Union[str, List[str]],
     max_sequence_length: int = 512,
     system_message: str = """You are an AI that reasons about image descriptions. You give structured responses focusing on object relationships, object
@@ -90,8 +90,8 @@ def prepare_text_ids(
 
 
 def encode_prompt(
-    text_encoder: Mistral3ForConditionalGeneration,
-    tokenizer: AutoProcessor,
+    text_encoder: AutoModelForCausalLM,
+    tokenizer: AutoTokenizer,
     prompt: Union[str, List[str]],
     num_images_per_prompt: int = 1,
     prompt_embeds: Optional[torch.Tensor] = None,