mirror of
https://github.com/justUmen/Bjornulf_custom_nodes.git
synced 2026-03-21 12:42:11 -03:00
29 lines
1.4 KiB
Python
29 lines
1.4 KiB
Python
import ollama
|
|
from ollama import Client # pip install ollama
|
|
|
|
class ollamaLoader:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"user_prompt": ("STRING", {"multiline": True}),
|
|
# "selected_model": ((), {}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("ollama_response",)
|
|
FUNCTION = "connect_2_ollama"
|
|
# INPUT_NODE = True # Changed from OUTPUT_NODE to INPUT_NODE
|
|
CATEGORY = "Bjornulf"
|
|
|
|
# @classmethod
|
|
def connect_2_ollama(self, user_prompt):
|
|
keep_alive = 0
|
|
list_models=ollama.list() #{'models': [{'name': 'dolphin-llama3:latest', 'model': 'dolphin-llama3:latest', 'modified_at': '2024-04-24T06:56:57.498527412+02:00', 'size': 4661235994, 'digest': '613f068e29f863bb900e568f920401b42678efca873d7a7c87b0d6ef4945fadd', 'details': {'parent_model': '', 'format': 'gguf', 'family': 'llama', 'families': ['llama'], 'parameter_size': '8B', 'quantization_level': 'Q4_0'}}]}
|
|
print(list_models)
|
|
client = Client(host="http://127.0.0.1:11434")
|
|
response = client.generate(model="dolphin-llama3", system="I will give you an object, animal, person or landscape, just create details about it : colors, size, clothes, eyes and other physical details or features in 1 sentence.", prompt=user_prompt, keep_alive=str(keep_alive) + "m")
|
|
print("Ollama response : ", response['response'])
|
|
return (response['response'],)
|