-
Notifications
You must be signed in to change notification settings - Fork 0
/
NanocastOllamabasic.py
66 lines (55 loc) · 2.06 KB
/
NanocastOllamabasic.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
import os
import requests
from dotenv import load_dotenv
from comfyui.core.node import Node, Input, Output
# Load environment variables from a .env file
load_dotenv()
class NanocastOllamaBasic(Node):
def __init__(self):
super().__init__()
self.api_endpoint = "http://localhost:8000/api"
self.api_key = os.getenv("API_KEY")
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"text": ("STRING", {
"multiline": True,
"dynamicPrompts": False,
"default": "Enter text here"
}),
},
}
RETURN_TYPES = ("STRING", "STRING")
RETURN_NAMES = ("narration", "image_prompts")
FUNCTION = "parse_and_condense"
CATEGORY = "Text Processing"
def parse_and_condense(self, text):
prompt = (
f"Here is some news feed content: {text}\n"
"You will create a narrative summary/script of this article. 1 article runs from Title to the next title. "
"Highlight 15-30 prompts that will generate images relating to this story. "
"Split these two outputs into 'narration' and 'image_prompts'."
)
response = requests.post(
self.api_endpoint,
headers={"Authorization": f"Bearer {self.api_key}"},
json={"text": prompt, "task": "parse_and_condense"}
)
if response.status_code == 200:
result = response.json()
narration = result.get("narration", "Error: No narration output.")
image_prompts = result.get("image_prompts", "Error: No image prompts output.")
else:
narration = "Error: Unable to parse and condense text."
image_prompts = "Error: Unable to generate image prompts."
return narration, image_prompts
def register():
node = NanocastOllamaBasic()
ComfyUI.register_node(node)
NODE_CLASS_MAPPINGS = {
"NanocastOllamaBasic": NanocastOllamaBasic
}
NODE_DISPLAY_NAME_MAPPINGS = {
"NanocastOllamaBasic": "NanocastOllamaBasic"
}