Add files via upload

This commit is contained in:
AI Lab
2025-05-30 22:29:56 -07:00
committed by GitHub
parent c54a3a296b
commit 9974e83707
2 changed files with 236 additions and 42 deletions
+91
View File
@@ -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
View File
@@ -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)",
}