Upload joycaption.py

joycaption.py

@@ -30,11 +30,11 @@ BASE_DIR = Path(__file__).resolve().parent
 device = "cuda" if torch.cuda.is_available() else "cpu"
 HF_TOKEN = os.environ.get("HF_TOKEN", None)
 use_inference_client = False
-PIXTRAL_PATH = "mistral-community/pixtral-12b"
+PIXTRAL_PATHS = ["mistral-community/pixtral-12b"]
 
 llm_models = {
     "Orenguteng/Llama-3.1-8B-Lexi-Uncensored-V2": None,
-    #PIXTRAL_PATH: None,
+    #PIXTRAL_PATHS[0]: None,
     "bunnycore/LLama-3.1-8B-Matrix": None,
     "Sao10K/Llama-3.1-8B-Stheno-v3.4": None,
     "unsloth/Meta-Llama-3.1-8B-bnb-4bit": None,
@@ -150,7 +150,7 @@ def load_text_model(model_name: str=MODEL_PATH, gguf_file: Union[str, None]=None
         nf4_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4",
                                         bnb_4bit_use_double_quant=True, bnb_4bit_compute_dtype=torch.bfloat16)
         
-        if model_name == PIXTRAL_PATH:
+        if model_name in PIXTRAL_PATHS:
             from transformers import AutoProcessor, LlavaForConditionalGeneration
             if is_nf4:
                 text_model = LlavaForConditionalGeneration.from_pretrained(model_name, quantization_config=nf4_config, device_map=device, torch_dtype=torch.bfloat16).eval()
@@ -249,7 +249,7 @@ def stream_chat_mod(input_image: Image.Image, caption_type: str, caption_tone: s
     print(f"Prompt: {prompt_str}")
 
     # Pixtral
-    if model_name == PIXTRAL_PATH:
+    if model_name in PIXTRAL_PATHS:
         input_images = [input_image]
         inputs = image_adapter(text=prompt_str, images=input_images, return_tensors="pt").to(device)
         generate_ids = text_model.generate(**inputs, max_new_tokens=max_new_tokens)