Add files via upload
This commit is contained in:
@@ -0,0 +1,91 @@
|
||||
import os
|
||||
import io
|
||||
import sys
|
||||
from pydub import AudioSegment
|
||||
from pydub.playback import play
|
||||
import torch
|
||||
import numpy as np
|
||||
from scipy.io import wavfile
|
||||
|
||||
class Everything(str):
|
||||
def __ne__(self, __value: object) -> bool:
|
||||
return False
|
||||
|
||||
class ailab_PlayAudio:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"optional": {
|
||||
"AUDIO": ("AUDIO", {"forceInput": True}),
|
||||
"audio_path": ("STRING", {"default": ""}),
|
||||
"autoplay": ("BOOLEAN", {"default": True, "label": "Auto Play"})
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("AUDIO", "STRING",)
|
||||
RETURN_NAMES = ("AUDIO", "AUDIO_PATH",)
|
||||
FUNCTION = "execute"
|
||||
CATEGORY = "🧪AILab/🌸Pollinations"
|
||||
OUTPUT_NODE = True
|
||||
ALWAYS_CHANGED = True # This tells ComfyUI to always execute this node
|
||||
|
||||
def play_audio(self, AUDIO=None, audio_path=None, autoplay=True):
|
||||
try:
|
||||
sound = None
|
||||
path = ""
|
||||
|
||||
if AUDIO is not None:
|
||||
if isinstance(AUDIO, dict) and 'waveform' in AUDIO:
|
||||
waveform = AUDIO['waveform']
|
||||
sample_rate = AUDIO.get('sample_rate', 44100)
|
||||
|
||||
if isinstance(waveform, torch.Tensor):
|
||||
waveform = waveform.cpu().numpy()
|
||||
|
||||
if waveform.dtype.kind == 'f':
|
||||
waveform = (waveform * 32767).astype(np.int16)
|
||||
|
||||
temp_wav = io.BytesIO()
|
||||
wavfile.write(temp_wav, sample_rate, waveform)
|
||||
temp_wav.seek(0)
|
||||
sound = AudioSegment.from_wav(temp_wav)
|
||||
path = ""
|
||||
|
||||
elif isinstance(AUDIO, AudioSegment):
|
||||
sound = AUDIO
|
||||
path = ""
|
||||
else:
|
||||
raise ValueError(f"Unsupported AUDIO type: {type(AUDIO)}")
|
||||
|
||||
elif audio_path and os.path.exists(audio_path):
|
||||
sound = AudioSegment.from_file(audio_path)
|
||||
path = audio_path
|
||||
|
||||
if autoplay and sound is not None:
|
||||
if sys.platform.startswith('win'):
|
||||
wav_io = io.BytesIO()
|
||||
sound.export(wav_io, format='wav')
|
||||
wav_data = wav_io.getvalue()
|
||||
import winsound
|
||||
winsound.PlaySound(wav_data, winsound.SND_MEMORY)
|
||||
else:
|
||||
play(sound)
|
||||
|
||||
return sound, path
|
||||
|
||||
except Exception as e:
|
||||
import traceback
|
||||
print(traceback.format_exc())
|
||||
return None, ""
|
||||
|
||||
def execute(self, AUDIO=None, audio_path=None, autoplay=True):
|
||||
sound, path = self.play_audio(AUDIO, audio_path, autoplay)
|
||||
return (sound, path)
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"ailab_PlayAudio": ailab_PlayAudio,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"ailab_PlayAudio": "Play Audio 🔊",
|
||||
}
|
||||
+145
-42
@@ -3,17 +3,23 @@ import sys
|
||||
import json
|
||||
import numpy as np
|
||||
import torch
|
||||
import torchaudio
|
||||
from PIL import Image
|
||||
import requests
|
||||
import tempfile
|
||||
import time
|
||||
from urllib.parse import quote, unquote
|
||||
from pydub import AudioSegment
|
||||
import io
|
||||
import folder_paths
|
||||
|
||||
DEFAULT_IMAGE_MODELS = ["flux", "flux-pro", "flux-realism", "flux-anime", "flux-3d", "flux-cablyai", "turbo"]
|
||||
DEFAULT_TEXT_MODELS = ["openai", "gpt-4", "gpt-3.5-turbo"]
|
||||
# Updated based on https://image.pollinations.ai/models
|
||||
DEFAULT_IMAGE_MODELS = ["flux", "turbo"]
|
||||
# Top models from https://text.pollinations.ai/models (first few models)
|
||||
DEFAULT_TEXT_MODELS = ["openai", "openai-fast", "openai-large", "qwen-coder", "llama", "mistral"]
|
||||
|
||||
MODELS_CACHE = {"models": [], "last_update": 0}
|
||||
TEXT_MODELS_CACHE = {"models": [], "last_update": 0}
|
||||
TEXT_MODELS_CACHE = {"model_info": [], "last_update": 0}
|
||||
|
||||
def get_available_models():
|
||||
"""Get available image models from API with caching"""
|
||||
@@ -41,25 +47,40 @@ def get_text_models():
|
||||
"""Get available text models from API with caching"""
|
||||
current_time = time.time()
|
||||
|
||||
if current_time - TEXT_MODELS_CACHE["last_update"] > 3600 or not TEXT_MODELS_CACHE["models"]:
|
||||
if current_time - TEXT_MODELS_CACHE["last_update"] > 3600 or not TEXT_MODELS_CACHE["model_info"]:
|
||||
try:
|
||||
response = requests.get("https://text.pollinations.ai/models", timeout=10)
|
||||
if response.status_code == 200:
|
||||
models_data = response.json()
|
||||
if models_data and len(models_data) > 0:
|
||||
# Extract only model names from response
|
||||
model_names = [model["name"] for model in models_data]
|
||||
TEXT_MODELS_CACHE["models"] = model_names
|
||||
# Store model info as {name, description} pairs
|
||||
model_info = []
|
||||
for model in models_data:
|
||||
model_info.append({
|
||||
'name': model['name'],
|
||||
'description': model.get('description', f"Default {model['name']} model")
|
||||
})
|
||||
TEXT_MODELS_CACHE["model_info"] = model_info
|
||||
else:
|
||||
TEXT_MODELS_CACHE["models"] = DEFAULT_TEXT_MODELS
|
||||
TEXT_MODELS_CACHE["model_info"] = [
|
||||
{'name': name, 'description': f"Default {name} model"}
|
||||
for name in DEFAULT_TEXT_MODELS
|
||||
]
|
||||
TEXT_MODELS_CACHE["last_update"] = current_time
|
||||
else:
|
||||
TEXT_MODELS_CACHE["models"] = DEFAULT_TEXT_MODELS
|
||||
TEXT_MODELS_CACHE["model_info"] = [
|
||||
{'name': name, 'description': f"Default {name} model"}
|
||||
for name in DEFAULT_TEXT_MODELS
|
||||
]
|
||||
except Exception as e:
|
||||
print(f"Error fetching text models: {e}")
|
||||
TEXT_MODELS_CACHE["models"] = DEFAULT_TEXT_MODELS
|
||||
TEXT_MODELS_CACHE["model_info"] = [
|
||||
{'name': name, 'description': f"Default {name} model"}
|
||||
for name in DEFAULT_TEXT_MODELS
|
||||
]
|
||||
|
||||
return TEXT_MODELS_CACHE["models"]
|
||||
# Return only descriptions for display
|
||||
return [model['description'] for model in TEXT_MODELS_CACHE["model_info"]]
|
||||
|
||||
class PollinationsImageGen:
|
||||
|
||||
@@ -88,9 +109,9 @@ class PollinationsImageGen:
|
||||
|
||||
RETURN_TYPES = ("IMAGE", "STRING", "STRING")
|
||||
RETURN_NAMES = ("images", "image_urls", "prompts")
|
||||
OUTPUT_IS_LIST = (True, False, False)
|
||||
OUTPUT_IS_LIST = (True, True, False)
|
||||
FUNCTION = "generate"
|
||||
CATEGORY = "🧪AILab/Pollinations"
|
||||
CATEGORY = "🧪AILab/🌸Pollinations"
|
||||
|
||||
def generate(self, prompt, model, width, height, batch_size=1, negative_prompt="", seed=0,
|
||||
enhance=True, nologo=True, private=True, safe=False):
|
||||
@@ -121,22 +142,17 @@ class PollinationsImageGen:
|
||||
enhance=True, nologo=True, private=True, safe=False):
|
||||
"""Generate a single image"""
|
||||
try:
|
||||
# Build base URL - using official API format from reference
|
||||
base_url = "https://image.pollinations.ai/prompt/"
|
||||
|
||||
# Build full prompt
|
||||
full_prompt = prompt
|
||||
if negative_prompt:
|
||||
full_prompt = f"{prompt} ### {negative_prompt}"
|
||||
|
||||
# URL encode the prompt
|
||||
encoded_prompt = quote(full_prompt)
|
||||
|
||||
# Build parameters
|
||||
params = {}
|
||||
params["model"] = model
|
||||
params["width"] = width
|
||||
params["height"] = height
|
||||
params = {
|
||||
"model": model,
|
||||
"width": width,
|
||||
"height": height,
|
||||
}
|
||||
|
||||
if seed and seed != 0:
|
||||
params["seed"] = seed
|
||||
@@ -149,20 +165,15 @@ class PollinationsImageGen:
|
||||
if safe:
|
||||
params["safe"] = "true"
|
||||
|
||||
# Build complete URL
|
||||
param_str = "&".join([f"{k}={v}" for k, v in params.items()])
|
||||
url = f"{base_url}{encoded_prompt}?{param_str}"
|
||||
|
||||
print(f"Generating image, URL: {url}")
|
||||
|
||||
# Download image
|
||||
response = requests.get(url, stream=True)
|
||||
response.raise_for_status()
|
||||
|
||||
# Get the final prompt used (if enhanced)
|
||||
final_prompt = full_prompt # Default to original prompt
|
||||
final_prompt = full_prompt
|
||||
|
||||
# Try to extract enhanced prompt from response URL
|
||||
try:
|
||||
image_url = response.url
|
||||
if "/prompt/" in image_url:
|
||||
@@ -174,7 +185,6 @@ class PollinationsImageGen:
|
||||
except Exception as ee:
|
||||
print(f"Error extracting enhanced prompt: {ee}")
|
||||
|
||||
# Save to temporary file
|
||||
temp_dir = tempfile.gettempdir()
|
||||
filename = f"pollinations_{int(time.time())}.png"
|
||||
image_path = os.path.join(temp_dir, filename)
|
||||
@@ -183,7 +193,6 @@ class PollinationsImageGen:
|
||||
for chunk in response.iter_content(chunk_size=8192):
|
||||
f.write(chunk)
|
||||
|
||||
# Load image
|
||||
image = Image.open(image_path)
|
||||
image_tensor = torch.from_numpy(np.array(image).astype(np.float32) / 255.0)[None,]
|
||||
|
||||
@@ -192,23 +201,23 @@ class PollinationsImageGen:
|
||||
except Exception as e:
|
||||
error_msg = f"Pollinations API error: {str(e)}"
|
||||
print(error_msg)
|
||||
# Return error message
|
||||
empty_image = torch.zeros(1, 512, 512, 3)
|
||||
return (empty_image, error_msg, prompt)
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, **kwargs):
|
||||
# Ensure a new image is generated each time
|
||||
return time.time()
|
||||
|
||||
class PollinationsTextGen:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
text_models = get_text_models()
|
||||
default_description = next((model['description'] for model in TEXT_MODELS_CACHE.get("model_info", [])
|
||||
if model['name'] == "openai"), text_models[0] if text_models else "Default openai model")
|
||||
return {
|
||||
"required": {
|
||||
"prompt": ("STRING", {"multiline": True, "placeholder": "Enter your text prompt..."}),
|
||||
"model": (text_models, {"default": "openai"}),
|
||||
"model": (text_models, {"default": default_description}),
|
||||
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
||||
},
|
||||
"optional": {
|
||||
@@ -219,19 +228,29 @@ class PollinationsTextGen:
|
||||
RETURN_TYPES = ("STRING",)
|
||||
RETURN_NAMES = ("generated_text",)
|
||||
FUNCTION = "generate_text"
|
||||
CATEGORY = "🧪AILab/Pollinations"
|
||||
CATEGORY = "🧪AILab/🌸Pollinations"
|
||||
|
||||
def generate_text(self, prompt, model, seed, private=True):
|
||||
try:
|
||||
# Build URL with parameters
|
||||
# Find the model name that matches this description
|
||||
model_name = None
|
||||
for model_info in TEXT_MODELS_CACHE["model_info"]:
|
||||
if model_info['description'] == model:
|
||||
model_name = model_info['name']
|
||||
break
|
||||
|
||||
if not model_name:
|
||||
model_name = "openai" # fallback
|
||||
|
||||
params = {
|
||||
"model": model,
|
||||
"model": model_name,
|
||||
"seed": seed,
|
||||
"private": str(private).lower()
|
||||
}
|
||||
param_str = "&".join([f"{k}={v}" for k, v in params.items()])
|
||||
url = f"https://text.pollinations.ai/{quote(prompt)}?{param_str}"
|
||||
|
||||
print(f"Generating Text, URL: {url}")
|
||||
response = requests.get(url)
|
||||
if response.status_code == 200:
|
||||
return (response.text,)
|
||||
@@ -240,15 +259,99 @@ class PollinationsTextGen:
|
||||
except Exception as e:
|
||||
return (f"Text generation failed: {str(e)}",)
|
||||
|
||||
# Adding Text-to-Speech node
|
||||
class PollinationsTextToSpeech:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"text": ("STRING", {"multiline": True, "placeholder": "Enter text to convert to speech..."}),
|
||||
"voice": (["nova", "alloy", "echo", "fable", "onyx", "shimmer"], {"default": "nova"}),
|
||||
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
||||
},
|
||||
"optional": {
|
||||
"private": ("BOOLEAN", {"default": True, "tooltip": "Keep the generation private"})
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("AUDIO", "STRING",)
|
||||
RETURN_NAMES = ("audio", "audio_path",)
|
||||
FUNCTION = "generate_speech"
|
||||
CATEGORY = "🧪AILab/🌸Pollinations"
|
||||
|
||||
def generate_speech(self, text, voice, seed, private=True):
|
||||
try:
|
||||
params = {
|
||||
"model": "openai-audio",
|
||||
"voice": voice,
|
||||
"seed": seed,
|
||||
"private": str(private).lower()
|
||||
}
|
||||
param_str = "&".join([f"{k}={v}" for k, v in params.items()])
|
||||
url = f"https://text.pollinations.ai/{quote(text)}?{param_str}"
|
||||
|
||||
print(f"Generating Speech, URL: {url}")
|
||||
response = requests.get(url, stream=True)
|
||||
|
||||
if response.status_code == 200:
|
||||
# Get ComfyUI's temp directory
|
||||
temp_dir = os.path.join(folder_paths.get_output_directory(), "pollinations_temp")
|
||||
os.makedirs(temp_dir, exist_ok=True)
|
||||
|
||||
# Generate unique filename
|
||||
timestamp = int(time.time())
|
||||
mp3_filename = f"pollinations_speech_{timestamp}.mp3"
|
||||
mp3_path = os.path.join(temp_dir, mp3_filename)
|
||||
|
||||
# Save MP3 file
|
||||
with open(mp3_path, 'wb') as f:
|
||||
for chunk in response.iter_content(chunk_size=8192):
|
||||
f.write(chunk)
|
||||
|
||||
# Load and process audio
|
||||
waveform, sample_rate = torchaudio.load(mp3_path)
|
||||
|
||||
# Ensure mono audio (take mean if stereo)
|
||||
if waveform.shape[0] > 1:
|
||||
waveform = waveform.mean(dim=0, keepdim=True)
|
||||
|
||||
# Add batch dimension if needed
|
||||
if waveform.dim() == 2:
|
||||
waveform = waveform.unsqueeze(0)
|
||||
|
||||
# Normalize audio
|
||||
if waveform.numel() > 0:
|
||||
max_val = waveform.abs().max()
|
||||
if max_val > 0:
|
||||
waveform = waveform / max_val
|
||||
|
||||
# Return audio in ComfyUI format
|
||||
audio_dict = {
|
||||
"waveform": waveform,
|
||||
"sample_rate": sample_rate
|
||||
}
|
||||
|
||||
return (audio_dict, mp3_path)
|
||||
else:
|
||||
print(f"Error generating speech: {response.status_code}")
|
||||
return ({"waveform": torch.zeros(1, 1, 16000), "sample_rate": 16000}, "")
|
||||
except Exception as e:
|
||||
error_msg = f"Speech generation failed: {str(e)}"
|
||||
print(error_msg)
|
||||
return ({"waveform": torch.zeros(1, 1, 16000), "sample_rate": 16000}, "")
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, **kwargs):
|
||||
return float("NaN")
|
||||
|
||||
# Register nodes
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"PollinationsImageGen": PollinationsImageGen,
|
||||
"PollinationsTextGen": PollinationsTextGen,
|
||||
"PollinationsTextToSpeech": PollinationsTextToSpeech,
|
||||
}
|
||||
|
||||
# UI display name
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"PollinationsImageGen": "Pollinations Image Gen 🖼️",
|
||||
"PollinationsTextGen": "Pollinations Text Gen 📝",
|
||||
}
|
||||
"PollinationsImageGen": "Image Gen 🖼️ (Pollinations)",
|
||||
"PollinationsTextGen": "Text Gen 📝 (Pollinations)",
|
||||
"PollinationsTextToSpeech": "Text To Speech Chat 🔊 (Pollinations)",
|
||||
}
|
||||
Reference in New Issue
Block a user