Files
2025-10-21 15:56:23 -04:00

597 lines
24 KiB
Python

import requests
import torch
import replicate
import os
import numpy as np
from PIL import Image
import io
import base64
def create_error_image():
"""Crear una imagen de error roja"""
blank = Image.new('RGB', (512, 512), color='red')
blank.putpixel((256, 256), (255, 0, 0))
tensor = torch.from_numpy(np.array(blank).astype(np.float32) / 255.0)
return tensor[None,]
class ReplicateQwenImageEdit:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"images": ("IMAGE",),
"prompt": ("STRING", {"multiline": True, "default": ""}),
"api_token": ("STRING", {"default": ""}),
},
"optional": {
"seed": ("INT", {"default": -1}),
"go_fast": ("BOOLEAN", {"default": True}),
"aspect_ratio": (["match_input_image", "1:1", "16:9", "9:16", "4:3", "3:4"], {"default": "match_input_image"}),
"output_format": (["webp", "jpg", "png"], {"default": "webp"}),
"output_quality": ("INT", {"default": 95, "min": 0, "max": 100}),
"disable_safety_checker": ("BOOLEAN", {"default": False}),
}
}
RETURN_TYPES = ("IMAGE", "STRING")
RETURN_NAMES = ("image", "output_url")
FUNCTION = "generate"
CATEGORY = "Replicate/Qwen"
def generate(self, **kwargs):
try:
# Obtener el token de API
api_token = self.get_api_token(kwargs['api_token'])
# Configurar el cliente de Replicate
client = replicate.Client(api_token=api_token)
# Procesar las imágenes (puede ser 1-5 imágenes)
images_tensor = kwargs['images']
# Verificar que tenemos al menos 1 imagen
if images_tensor is None or len(images_tensor) == 0:
raise Exception("Se requiere al menos 1 imagen")
# Verificar que no excedemos el máximo de 5 imágenes
if len(images_tensor) > 5:
raise Exception("Máximo 5 imágenes permitidas")
# Convertir todas las imágenes a URLs
image_urls = []
for i in range(len(images_tensor)):
image_url = self.image_to_url(images_tensor[i])
image_urls.append(image_url)
# Preparar los parámetros de entrada
input_params = {
"image": image_urls,
"prompt": kwargs['prompt'],
"go_fast": kwargs.get('go_fast', True),
"aspect_ratio": kwargs.get('aspect_ratio', "match_input_image"),
"output_format": kwargs.get('output_format', "webp"),
"output_quality": kwargs.get('output_quality', 95),
"disable_safety_checker": kwargs.get('disable_safety_checker', False),
}
# Agregar seed solo si no es -1
if kwargs.get('seed', -1) != -1:
input_params["seed"] = kwargs['seed']
# Ejecutar el modelo
output = client.run(
"qwen/qwen-image-edit-plus",
input=input_params
)
# Procesar la salida
if output and len(output) > 0:
# Obtener la URL de la imagen generada
# La salida de Replicate ya es una URL de cadena directamente
output_url = str(output[0])
# Descargar y convertir la imagen
image_tensor = self.download_and_convert_image(output_url)
return (image_tensor, output_url)
else:
raise Exception("No se recibió ninguna imagen de la API")
except Exception as e:
print(f"Error en Replicate Qwen Image Edit: {str(e)}")
return (create_error_image(), f"Error: {str(e)}")
def get_api_token(self, node_token):
"""Obtener el token de API desde el nodo o variables de entorno"""
if node_token and node_token.strip():
return node_token.strip()
# Intentar obtener desde variable de entorno
env_token = os.getenv('REPLICATE_API_TOKEN')
if env_token:
return env_token
raise ValueError("No se encontró un token de API válido. Configure REPLICATE_API_TOKEN en las variables de entorno o en el nodo.")
def image_to_url(self, image_tensor):
"""Convertir tensor de imagen a URL temporal o base64"""
try:
# Convertir tensor a PIL Image
if hasattr(image_tensor, 'numpy'):
image_array = (image_tensor.numpy().squeeze() * 255).astype(np.uint8)
else:
# Si ya es un array de numpy
image_array = (image_tensor.squeeze() * 255).astype(np.uint8)
if len(image_array.shape) == 3:
image = Image.fromarray(image_array)
else:
image = Image.fromarray(image_array, mode='L')
# Convertir a base64 data URL
buffered = io.BytesIO()
image.save(buffered, format='PNG')
img_base64 = base64.b64encode(buffered.getvalue()).decode('utf-8')
return f"data:image/png;base64,{img_base64}"
except Exception as e:
print(f"Error convirtiendo imagen: {str(e)}")
print(f"Tipo de tensor: {type(image_tensor)}")
print(f"Forma del tensor: {image_tensor.shape if hasattr(image_tensor, 'shape') else 'No shape'}")
raise Exception(f"Error procesando imagen: {str(e)}")
def download_and_convert_image(self, image_url):
"""Descargar imagen desde URL y convertir a tensor"""
try:
# Descargar la imagen
response = requests.get(image_url, timeout=30)
response.raise_for_status()
# Convertir a PIL Image
image = Image.open(io.BytesIO(response.content))
# Convertir a RGB si es necesario
if image.mode != 'RGB':
image = image.convert('RGB')
# Convertir a tensor de ComfyUI
img_array = np.array(image).astype(np.float32) / 255.0
return torch.from_numpy(img_array)[None,]
except Exception as e:
print(f"Error descargando imagen: {str(e)}")
raise Exception(f"Error descargando imagen: {str(e)}")
class ReplicateQwenImageEditSingle:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"image": ("IMAGE",),
"prompt": ("STRING", {"multiline": True, "default": ""}),
"api_token": ("STRING", {"default": ""}),
},
"optional": {
"seed": ("INT", {"default": -1}),
"go_fast": ("BOOLEAN", {"default": True}),
"aspect_ratio": (["match_input_image", "1:1", "16:9", "9:16", "4:3", "3:4"], {"default": "match_input_image"}),
"output_format": (["webp", "jpg", "png"], {"default": "webp"}),
"output_quality": ("INT", {"default": 95, "min": 0, "max": 100}),
"disable_safety_checker": ("BOOLEAN", {"default": False}),
}
}
RETURN_TYPES = ("IMAGE", "STRING")
RETURN_NAMES = ("image", "output_url")
FUNCTION = "generate"
CATEGORY = "Replicate/Qwen"
def generate(self, **kwargs):
try:
# Obtener el token de API
api_token = self.get_api_token(kwargs['api_token'])
# Configurar el cliente de Replicate
client = replicate.Client(api_token=api_token)
# Procesar la imagen (solo una imagen)
image_tensor = kwargs['image']
# Verificar que tenemos una imagen
if image_tensor is None or len(image_tensor) == 0:
raise Exception("Se requiere al menos 1 imagen")
# Convertir la imagen a URL
image_url = self.image_to_url(image_tensor[0]) # Tomar la primera imagen
# Preparar los parámetros de entrada
input_params = {
"image": image_url,
"prompt": kwargs['prompt'],
"go_fast": kwargs.get('go_fast', True),
"aspect_ratio": kwargs.get('aspect_ratio', "match_input_image"),
"output_format": kwargs.get('output_format', "webp"),
"output_quality": kwargs.get('output_quality', 95),
"disable_safety_checker": kwargs.get('disable_safety_checker', False),
}
# Agregar seed solo si no es -1
if kwargs.get('seed', -1) != -1:
input_params["seed"] = kwargs['seed']
# Ejecutar el modelo
output = client.run(
"qwen/qwen-image-edit",
input=input_params
)
# Procesar la salida
if output and len(output) > 0:
# Obtener la URL de la imagen generada
output_url = str(output[0])
# Descargar y convertir la imagen
image_tensor = self.download_and_convert_image(output_url)
return (image_tensor, output_url)
else:
raise Exception("No se recibió ninguna imagen de la API")
except Exception as e:
print(f"Error en Replicate Qwen Image Edit Single: {str(e)}")
return (create_error_image(), f"Error: {str(e)}")
def get_api_token(self, node_token):
"""Obtener el token de API desde el nodo o variables de entorno"""
if node_token and node_token.strip():
return node_token.strip()
# Intentar obtener desde variable de entorno
env_token = os.getenv('REPLICATE_API_TOKEN')
if env_token:
return env_token
raise ValueError("No se encontró un token de API válido. Configure REPLICATE_API_TOKEN en las variables de entorno o en el nodo.")
def image_to_url(self, image_tensor):
"""Convertir tensor de imagen a URL temporal o base64"""
try:
# Convertir tensor a PIL Image
if hasattr(image_tensor, 'numpy'):
image_array = (image_tensor.numpy().squeeze() * 255).astype(np.uint8)
else:
# Si ya es un array de numpy
image_array = (image_tensor.squeeze() * 255).astype(np.uint8)
if len(image_array.shape) == 3:
image = Image.fromarray(image_array)
else:
image = Image.fromarray(image_array, mode='L')
# Convertir a base64 data URL
buffered = io.BytesIO()
image.save(buffered, format='PNG')
img_base64 = base64.b64encode(buffered.getvalue()).decode('utf-8')
return f"data:image/png;base64,{img_base64}"
except Exception as e:
print(f"Error convirtiendo imagen: {str(e)}")
print(f"Tipo de tensor: {type(image_tensor)}")
print(f"Forma del tensor: {image_tensor.shape if hasattr(image_tensor, 'shape') else 'No shape'}")
raise Exception(f"Error procesando imagen: {str(e)}")
def download_and_convert_image(self, image_url):
"""Descargar imagen desde URL y convertir a tensor"""
try:
# Descargar la imagen
response = requests.get(image_url, timeout=30)
response.raise_for_status()
# Convertir a PIL Image
image = Image.open(io.BytesIO(response.content))
# Convertir a RGB si es necesario
if image.mode != 'RGB':
image = image.convert('RGB')
# Convertir a tensor de ComfyUI
img_array = np.array(image).astype(np.float32) / 255.0
return torch.from_numpy(img_array)[None,]
except Exception as e:
print(f"Error descargando imagen: {str(e)}")
raise Exception(f"Error descargando imagen: {str(e)}")
class ReplicateSeedEdit:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"image": ("IMAGE",),
"prompt": ("STRING", {"multiline": True, "default": ""}),
"api_token": ("STRING", {"default": ""}),
},
"optional": {
"seed": ("INT", {"default": -1}),
"guidance_scale": ("FLOAT", {"default": 5.5, "min": 1.0, "max": 10.0}),
}
}
RETURN_TYPES = ("IMAGE", "STRING")
RETURN_NAMES = ("image", "output_url")
FUNCTION = "generate"
CATEGORY = "Replicate/SeedEdit"
def generate(self, **kwargs):
try:
# Obtener el token de API
api_token = self.get_api_token(kwargs['api_token'])
# Configurar el cliente de Replicate
client = replicate.Client(api_token=api_token)
# Procesar la imagen (solo una imagen)
image_tensor = kwargs['image']
# Verificar que tenemos una imagen
if image_tensor is None or len(image_tensor) == 0:
raise Exception("Se requiere al menos 1 imagen")
# Convertir la imagen a URL
image_url = self.image_to_url(image_tensor[0]) # Tomar la primera imagen
# Preparar los parámetros de entrada
input_params = {
"image": image_url,
"prompt": kwargs['prompt'],
"guidance_scale": kwargs.get('guidance_scale', 5.5),
}
# Agregar seed solo si no es -1
if kwargs.get('seed', -1) != -1:
input_params["seed"] = kwargs['seed']
# Ejecutar el modelo
output = client.run(
"bytedance/seededit-3.0",
input=input_params
)
# Procesar la salida
if output:
# Obtener la URL de la imagen generada
# Para SeedEdit, la salida es directamente una URL de cadena
output_url = str(output)
# Descargar y convertir la imagen
image_tensor = self.download_and_convert_image(output_url)
return (image_tensor, output_url)
else:
raise Exception("No se recibió ninguna imagen de la API")
except Exception as e:
print(f"Error en Replicate SeedEdit: {str(e)}")
return (create_error_image(), f"Error: {str(e)}")
def get_api_token(self, node_token):
"""Obtener el token de API desde el nodo o variables de entorno"""
if node_token and node_token.strip():
return node_token.strip()
# Intentar obtener desde variable de entorno
env_token = os.getenv('REPLICATE_API_TOKEN')
if env_token:
return env_token
raise ValueError("No se encontró un token de API válido. Configure REPLICATE_API_TOKEN en las variables de entorno o en el nodo.")
def image_to_url(self, image_tensor):
"""Convertir tensor de imagen a URL temporal o base64"""
try:
# Convertir tensor a PIL Image
if hasattr(image_tensor, 'numpy'):
image_array = (image_tensor.numpy().squeeze() * 255).astype(np.uint8)
else:
# Si ya es un array de numpy
image_array = (image_tensor.squeeze() * 255).astype(np.uint8)
if len(image_array.shape) == 3:
image = Image.fromarray(image_array)
else:
image = Image.fromarray(image_array, mode='L')
# Convertir a base64 data URL
buffered = io.BytesIO()
image.save(buffered, format='PNG')
img_base64 = base64.b64encode(buffered.getvalue()).decode('utf-8')
return f"data:image/png;base64,{img_base64}"
except Exception as e:
print(f"Error convirtiendo imagen: {str(e)}")
print(f"Tipo de tensor: {type(image_tensor)}")
print(f"Forma del tensor: {image_tensor.shape if hasattr(image_tensor, 'shape') else 'No shape'}")
raise Exception(f"Error procesando imagen: {str(e)}")
def download_and_convert_image(self, image_url):
"""Descargar imagen desde URL y convertir a tensor"""
try:
# Descargar la imagen
response = requests.get(image_url, timeout=30)
response.raise_for_status()
# Convertir a PIL Image
image = Image.open(io.BytesIO(response.content))
# Convertir a RGB si es necesario
if image.mode != 'RGB':
image = image.convert('RGB')
# Convertir a tensor de ComfyUI
img_array = np.array(image).astype(np.float32) / 255.0
return torch.from_numpy(img_array)[None,]
except Exception as e:
print(f"Error descargando imagen: {str(e)}")
raise Exception(f"Error descargando imagen: {str(e)}")
class ReplicateSeedream4:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"prompt": ("STRING", {"multiline": True, "default": ""}),
"api_token": ("STRING", {"default": ""}),
},
"optional": {
"image_input": ("IMAGE",),
"size": (["1K", "2K", "4K", "custom"], {"default": "2K"}),
"width": ("INT", {"default": 2048, "min": 1024, "max": 4096}),
"height": ("INT", {"default": 2048, "min": 1024, "max": 4096}),
"max_images": ("INT", {"default": 1, "min": 1, "max": 15}),
"aspect_ratio": (["match_input_image", "1:1", "16:9", "9:16", "4:3", "3:4"], {"default": "match_input_image"}),
"sequential_image_generation": (["disabled", "auto"], {"default": "disabled"}),
}
}
RETURN_TYPES = ("IMAGE", "STRING")
RETURN_NAMES = ("image", "output_url")
FUNCTION = "generate"
CATEGORY = "Replicate/Seedream"
def generate(self, **kwargs):
try:
# Obtener el token de API
api_token = self.get_api_token(kwargs['api_token'])
# Configurar el cliente de Replicate
client = replicate.Client(api_token=api_token)
# Preparar los parámetros de entrada
input_params = {
"prompt": kwargs['prompt'],
"size": kwargs.get('size', "2K"),
"max_images": kwargs.get('max_images', 1),
"aspect_ratio": kwargs.get('aspect_ratio', "match_input_image"),
"sequential_image_generation": kwargs.get('sequential_image_generation', "disabled"),
}
# Agregar dimensiones personalizadas solo si size es 'custom'
if kwargs.get('size') == 'custom':
input_params['width'] = kwargs.get('width', 2048)
input_params['height'] = kwargs.get('height', 2048)
# Procesar imagen de entrada si se proporciona
if 'image_input' in kwargs and kwargs['image_input'] is not None:
image_tensor = kwargs['image_input']
# Verificar que tenemos una imagen
if image_tensor is not None and len(image_tensor) > 0:
# Convertir todas las imágenes a URLs
image_urls = []
for i in range(len(image_tensor)):
image_url = self.image_to_url(image_tensor[i])
image_urls.append(image_url)
input_params['image_input'] = image_urls
# Ejecutar el modelo
output = client.run(
"bytedance/seedream-4",
input=input_params
)
# Procesar la salida
if output and len(output) > 0:
# Obtener la URL de la imagen generada (tomar la primera si hay múltiples)
output_url = str(output[0])
# Descargar y convertir la imagen
image_tensor = self.download_and_convert_image(output_url)
return (image_tensor, output_url)
else:
raise Exception("No se recibió ninguna imagen de la API")
except Exception as e:
print(f"Error en Replicate Seedream4: {str(e)}")
return (create_error_image(), f"Error: {str(e)}")
def get_api_token(self, node_token):
"""Obtener el token de API desde el nodo o variables de entorno"""
if node_token and node_token.strip():
return node_token.strip()
# Intentar obtener desde variable de entorno
env_token = os.getenv('REPLICATE_API_TOKEN')
if env_token:
return env_token
raise ValueError("No se encontró un token de API válido. Configure REPLICATE_API_TOKEN en las variables de entorno o en el nodo.")
def image_to_url(self, image_tensor):
"""Convertir tensor de imagen a URL temporal o base64"""
try:
# Convertir tensor a PIL Image
if hasattr(image_tensor, 'numpy'):
image_array = (image_tensor.numpy().squeeze() * 255).astype(np.uint8)
else:
# Si ya es un array de numpy
image_array = (image_tensor.squeeze() * 255).astype(np.uint8)
if len(image_array.shape) == 3:
image = Image.fromarray(image_array)
else:
image = Image.fromarray(image_array, mode='L')
# Convertir a base64 data URL
buffered = io.BytesIO()
image.save(buffered, format='PNG')
img_base64 = base64.b64encode(buffered.getvalue()).decode('utf-8')
return f"data:image/png;base64,{img_base64}"
except Exception as e:
print(f"Error convirtiendo imagen: {str(e)}")
print(f"Tipo de tensor: {type(image_tensor)}")
print(f"Forma del tensor: {image_tensor.shape if hasattr(image_tensor, 'shape') else 'No shape'}")
raise Exception(f"Error procesando imagen: {str(e)}")
def download_and_convert_image(self, image_url):
"""Descargar imagen desde URL y convertir a tensor"""
try:
# Descargar la imagen
response = requests.get(image_url, timeout=30)
response.raise_for_status()
# Convertir a PIL Image
image = Image.open(io.BytesIO(response.content))
# Convertir a RGB si es necesario
if image.mode != 'RGB':
image = image.convert('RGB')
# Convertir a tensor de ComfyUI
img_array = np.array(image).astype(np.float32) / 255.0
return torch.from_numpy(img_array)[None,]
except Exception as e:
print(f"Error descargando imagen: {str(e)}")
raise Exception(f"Error descargando imagen: {str(e)}")
NODE_CLASS_MAPPINGS = {
"Replicate Qwen Image Edit Plus": ReplicateQwenImageEdit,
"Replicate Qwen Image Edit": ReplicateQwenImageEditSingle,
"Replicate SeedEdit": ReplicateSeedEdit,
"Replicate Seedream4": ReplicateSeedream4,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"Replicate Qwen Image Edit Plus": "Qwen Image Edit Plus (Multi-Image)",
"Replicate Qwen Image Edit": "Qwen Image Edit (Single Image)",
"Replicate SeedEdit": "SeedEdit 3.0 (ByteDance)",
"Replicate Seedream4": "Seedream 4.0 (ByteDance)",
}