From e77065dd7343dc4f952d45b293eb35ac83ccc1cd Mon Sep 17 00:00:00 2001
From: Celaleddin <celoron@gmail.com>
Date: Wed, 28 Aug 2024 14:37:17 +0300
Subject: [PATCH] initial commit

---
 .gitignore             | 87 ++++++++++++++++++++++++++++++++++++++++++
 VisualQueryTemplate.py | 64 +++++++++++++++++++++++++++++++
 __init__.py            | 10 +++++
 requirements.txt       |  3 ++
 4 files changed, 164 insertions(+)
 create mode 100644 .gitignore
 create mode 100644 VisualQueryTemplate.py
 create mode 100644 __init__.py
 create mode 100644 requirements.txt

diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..86a3f1f
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,87 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+*.egg-info/
+.installed.cfg
+*.egg
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# PyCharm
+.idea/
+
+# VS Code
+.vscode/
+
+# macOS
+.DS_Store
+
+# Linux
+*.swp
+
+# Windows
+Thumbs.db
+ehthumbs.db
+Desktop.ini
+
+# ComfyUI specific
+config.json
+logs/
+cache/
+
+# Custom node specific
+node_modules/
+*.log
+*.tmp
\ No newline at end of file
diff --git a/VisualQueryTemplate.py b/VisualQueryTemplate.py
new file mode 100644
index 0000000..231f548
--- /dev/null
+++ b/VisualQueryTemplate.py
@@ -0,0 +1,64 @@
+import re
+from PIL import Image
+import numpy as np
+import time
+from transformers import pipeline
+import torch
+
+def tensor2pil(image):
+    return Image.fromarray(np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8))
+
+class VisualQueryTemplateNode:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "images": ("IMAGE",),
+                "model": (["Salesforce/blip-vqa-base", "Salesforce/blip-vqa-capfilt-large", "dandelin/vilt-b32-finetuned-vqa", "microsoft/git-large-vqav2"], ),
+                "question": ("STRING", {"default": "{eye color} eyes, {hair style} {hair color} hair, {ethnicity} {gender}, {age number} years old, {facialhair}", "multiline": True, "dynamicPrompts": False}),
+            }
+        }
+
+    RETURN_TYPES = ("STRING",)
+    OUTPUT_IS_LIST = (True,)
+
+    FUNCTION = "vqa_image"
+    CATEGORY = "image"
+
+    def vqa_image(self, images, model, question):
+        start_time = time.time()
+
+        device = 0 if torch.cuda.is_available() else -1
+        vqa = pipeline(model=model, device=device)
+
+        answers = []
+
+        for image in images:
+            pil_image = tensor2pil(image).convert("RGB")
+
+            final_answer = question
+
+            matches = re.findall(r'\{([^}]*)\}', question)
+
+            for match in matches:
+
+                match_answers = vqa(question=match, image=pil_image)
+
+                print(match, match_answers)
+
+                match_answer = match_answers[0]["answer"]
+
+                final_answer = final_answer.replace("{"+match+"}", match_answer)
+
+            
+            answers.append(final_answer)
+
+        
+        end_time = time.time()
+        execution_time = end_time - start_time
+        print(f"Execution time: {execution_time} seconds")
+
+        return (answers,)
diff --git a/__init__.py b/__init__.py
new file mode 100644
index 0000000..f49a688
--- /dev/null
+++ b/__init__.py
@@ -0,0 +1,10 @@
+from .VisualQueryTemplate import VisualQueryTemplateNode
+
+NODE_CLASS_MAPPINGS = {
+    "VisualQueryTemplateNode": VisualQueryTemplateNode
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    # --- MAIN NODES ---
+    "VisualQueryTemplateNode": "Visual Query Template",
+}
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000..d707987
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1,3 @@
+torch
+numpy
+transformers
\ No newline at end of file