torch_dtype quantization is added for input gen

2024-08-26 20:26:52 +03:00
parent 00560c3c70
commit 5f359a8730
1 changed files with 4 additions and 1 deletions
@@ -148,7 +148,10 @@ class ImageCaptionNode:

        processor = transformers.AutoProcessor.from_pretrained(model_path)

-        inputs = processor(images=img, return_tensors="pt").to(dev)
+        try:
+            inputs = processor(images=img, return_tensors="pt").to(dev, req_torch_dtype)
+        except:
+            inputs = processor(images=img, return_tensors="pt").to(dev)

        out = model.generate(
            **inputs,