pip install torch transformers pillow Minimal Captioning Exampleimport osimport torchfrom PIL import Imagefrom transformers import AutoProcessor, Mistral3ForConditionalGeneration # =========================# Config# ========================= MODEL_PATH = "Felldude/Ministral-3-8B-Uncensored"FOLDER = "images"PROMPT = "Describe this image in detail."MAX_TOKENS = 512 VALID_EXTS = {".png", ".jpg", ".jpeg", ".webp"} # =========================# GPU setup# ========================= if not torch.cuda.is_available(): raise RuntimeError("CUDA GPU required") device = "cuda" dtype = ( torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16) # =========================# Load model# ========================= processor = AutoProcessor.from_pretrained( MODEL_PATH, trust_remote_code=True) model = Mistral3ForConditionalGeneration.from_pretrained( MODEL_PATH, torch_dtype=dtype, trust_remote_code=True, attn_implementation="sdpa").to(device) model.eval() # =========================# Caption function# ========================= def generate_caption(image): messages = [ { "role": "user", "content": [ {"type": "image", "image": image}, {"type": "text", "text": PROMPT}, ], } ] inputs = processor.apply_chat_template( messages, tokenize=True, add_generation_prompt=True, return_tensors="pt", return_dict=True, ) inputs = { k: v.to(device) for k, v in inputs.items() } with torch.inference_mode(): output = model.generate( **inputs, max_new_tokens=MAX_TOKENS, do_sample=False ) trimmed = [ o[len(i):] for i, o in zip(inputs["input_ids"], output) ] return processor.batch_decode( trimmed, skip_special_tokens=True )[0].strip() # =========================# Process folder# ========================= for filename in os.listdir(FOLDER): ext = os.path.splitext(filename)[1].lower() if ext not in VALID_EXTS: continue path = os.path.join(FOLDER, filename) print("Processing:", filename) try: image = Image.open(path).convert("RGB") caption = generate_caption(image) txt_path = os.path.splitext(path)[0] + ".txt" with open(txt_path, "w", encoding="utf-8") as f: f.write(caption) print(caption) except Exception as e: print("Failed:", e)