mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-01-11 23:00:51 +08:00
add args for ideogram nodes, add tests
This commit is contained in:
Benjamin Berman
parent
64fa54c3ad
commit
e6f9a6a552
@ -14,28 +14,39 @@ from comfy.utils import pil2tensor, tensor2pil
|
||||
from comfy_extras.constants.resolutions import IDEOGRAM_RESOLUTIONS
|
||||
from comfy_extras.nodes.nodes_mask import MaskToImage
|
||||
|
||||
# --- ENUMs and Constants ---
|
||||
|
||||
ASPECT_RATIOS = [(10, 6), (16, 10), (9, 16), (3, 2), (4, 3)]
|
||||
ASPECT_RATIO_ENUM = ["ASPECT_1_1"] + list(chain.from_iterable(
|
||||
[f"ASPECT_{a}_{b}", f"ASPECT_{b}_{a}"]
|
||||
for a, b in ASPECT_RATIOS
|
||||
))
|
||||
# New enum for v3 aspect ratios
|
||||
ASPECT_RATIO_V3_ENUM = ["disabled", "1x1", "10x16", "9x16", "3x4", "2x3", "16x10", "3x2", "4x3", "16x9"]
|
||||
V2_MODELS = ["V_2", "V_2_TURBO"]
|
||||
MODELS_ENUM = V2_MODELS + ["V_3"]
|
||||
AUTO_PROMPT_ENUM = ["AUTO", "ON", "OFF"]
|
||||
STYLES_ENUM = ["AUTO", "GENERAL", "REALISTIC", "DESIGN", "RENDER_3D", "ANIME"]
|
||||
RESOLUTION_ENUM = [f"RESOLUTION_{w}_{h}" for w, h in IDEOGRAM_RESOLUTIONS]
|
||||
# New enum for v3 rendering speed
|
||||
RENDERING_SPEED_ENUM = ["DEFAULT", "TURBO", "QUALITY"]
|
||||
|
||||
def to_v3_resolution(resolution:str) -> str:
|
||||
|
||||
# --- Helper Functions ---
|
||||
|
||||
def to_v3_resolution(resolution: str) -> str:
|
||||
return resolution[len("RESOLUTION_"):].replace("_", "x")
|
||||
|
||||
|
||||
def api_key_in_env_or_workflow(api_key_from_workflow: str):
|
||||
from comfy.cli_args import args
|
||||
if api_key_from_workflow is not None and "" != api_key_from_workflow.strip():
|
||||
return api_key_from_workflow
|
||||
|
||||
return os.environ.get("IDEOGRAM_API_KEY", args.ideogram_api_key)
|
||||
|
||||
|
||||
# --- Custom Nodes ---
|
||||
|
||||
class IdeogramGenerate(CustomNode):
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls) -> Dict[str, Any]:
|
||||
@ -43,7 +54,7 @@ class IdeogramGenerate(CustomNode):
|
||||
"required": {
|
||||
"prompt": ("STRING", {"multiline": True}),
|
||||
"resolution": (RESOLUTION_ENUM, {"default": RESOLUTION_ENUM[0]}),
|
||||
"model": (MODELS_ENUM, {"default": MODELS_ENUM[0]}),
|
||||
"model": (MODELS_ENUM, {"default": MODELS_ENUM[-1]}),
|
||||
"magic_prompt_option": (AUTO_PROMPT_ENUM, {"default": AUTO_PROMPT_ENUM[0]}),
|
||||
},
|
||||
"optional": {
|
||||
@ -52,6 +63,10 @@ class IdeogramGenerate(CustomNode):
|
||||
"num_images": ("INT", {"default": 1, "min": 1, "max": 8}),
|
||||
"seed": Seed,
|
||||
"style_type": (STYLES_ENUM, {}),
|
||||
# New v3 optional args
|
||||
"rendering_speed": (RENDERING_SPEED_ENUM, {"default": "DEFAULT"}),
|
||||
"aspect_ratio": (ASPECT_RATIO_V3_ENUM, {"default": "disabled"}),
|
||||
"style_reference_images": ("IMAGE",),
|
||||
}
|
||||
}
|
||||
|
||||
@ -60,56 +75,61 @@ class IdeogramGenerate(CustomNode):
|
||||
CATEGORY = "ideogram"
|
||||
|
||||
def generate(self, prompt: str, resolution: str, model: str, magic_prompt_option: str,
|
||||
api_key: str = "", negative_prompt: str = "", num_images: int = 1, seed: int = 0, style_type: str = "AUTO") -> Tuple[torch.Tensor]:
|
||||
api_key: str = "", negative_prompt: str = "", num_images: int = 1, seed: int = 0, style_type: str = "AUTO",
|
||||
rendering_speed: str = "DEFAULT", aspect_ratio: str = "disabled", style_reference_images: ImageBatch = None) -> Tuple[torch.Tensor]:
|
||||
api_key = api_key_in_env_or_workflow(api_key)
|
||||
headers = {"Api-Key": api_key, "Content-Type": "application/json"}
|
||||
|
||||
if model in V2_MODELS:
|
||||
headers = {"Api-Key": api_key, "Content-Type": "application/json"}
|
||||
payload = {
|
||||
"image_request": {
|
||||
"prompt": prompt,
|
||||
"resolution": resolution,
|
||||
"model": model,
|
||||
"magic_prompt_option": magic_prompt_option,
|
||||
"num_images": num_images,
|
||||
"prompt": prompt, "resolution": resolution, "model": model,
|
||||
"magic_prompt_option": magic_prompt_option, "num_images": num_images,
|
||||
"style_type": style_type,
|
||||
}
|
||||
}
|
||||
|
||||
if negative_prompt:
|
||||
payload["image_request"]["negative_prompt"] = negative_prompt
|
||||
if seed:
|
||||
payload["image_request"]["seed"] = seed
|
||||
|
||||
if negative_prompt: payload["image_request"]["negative_prompt"] = negative_prompt
|
||||
if seed: payload["image_request"]["seed"] = seed
|
||||
response = requests.post("https://api.ideogram.ai/generate", headers=headers, json=payload)
|
||||
|
||||
elif model == "V_3":
|
||||
payload = {
|
||||
"prompt": prompt,
|
||||
"resolution": to_v3_resolution(resolution),
|
||||
"model": model,
|
||||
"magic_prompt": magic_prompt_option,
|
||||
"num_images": num_images,
|
||||
"style_type": style_type,
|
||||
"prompt": prompt, "model": model, "magic_prompt": magic_prompt_option,
|
||||
"num_images": num_images, "style_type": style_type, "rendering_speed": rendering_speed,
|
||||
}
|
||||
if negative_prompt: payload["negative_prompt"] = negative_prompt
|
||||
if seed: payload["seed"] = seed
|
||||
|
||||
if negative_prompt:
|
||||
payload["negative_prompt"] = negative_prompt
|
||||
if seed:
|
||||
payload["seed"] = seed
|
||||
response = requests.post("https://api.ideogram.ai/v1/ideogram-v3/generate", headers=headers, json=payload)
|
||||
# Handle resolution vs aspect_ratio (aspect_ratio takes precedence)
|
||||
if aspect_ratio != "disabled":
|
||||
payload["aspect_ratio"] = aspect_ratio
|
||||
else:
|
||||
payload["resolution"] = to_v3_resolution(resolution)
|
||||
|
||||
headers = {"Api-Key": api_key}
|
||||
|
||||
# Use multipart/form-data if style references are provided
|
||||
if style_reference_images is not None:
|
||||
files = []
|
||||
for i, style_image in enumerate(style_reference_images):
|
||||
pil_image = tensor2pil(style_image)
|
||||
image_bytes = BytesIO()
|
||||
pil_image.save(image_bytes, format="PNG")
|
||||
files.append(("style_reference_images", (f"style_{i}.png", image_bytes.getvalue(), "image/png")))
|
||||
response = requests.post("https://api.ideogram.ai/v1/ideogram-v3/generate", headers=headers, data=payload, files=files)
|
||||
else:
|
||||
headers["Content-Type"] = "application/json"
|
||||
response = requests.post("https://api.ideogram.ai/v1/ideogram-v3/generate", headers=headers, json=payload)
|
||||
else:
|
||||
raise ValueError(f"Invalid model={model}")
|
||||
|
||||
response.raise_for_status()
|
||||
|
||||
images = []
|
||||
for item in response.json()["data"]:
|
||||
img_response = requests.get(item["url"])
|
||||
img_response.raise_for_status()
|
||||
|
||||
pil_image = Image.open(BytesIO(img_response.content))
|
||||
images.append(pil2tensor(pil_image))
|
||||
|
||||
return (torch.cat(images, dim=0),)
|
||||
|
||||
|
||||
@ -118,17 +138,17 @@ class IdeogramEdit(CustomNode):
|
||||
def INPUT_TYPES(cls) -> Dict[str, Any]:
|
||||
return {
|
||||
"required": {
|
||||
"images": ("IMAGE",),
|
||||
"masks": ("MASK",),
|
||||
"prompt": ("STRING", {"multiline": True}),
|
||||
"model": (MODELS_ENUM, {"default": MODELS_ENUM[0]}),
|
||||
"images": ("IMAGE",), "masks": ("MASK",), "prompt": ("STRING", {"multiline": True}),
|
||||
"model": (MODELS_ENUM, {"default": MODELS_ENUM[-1]}),
|
||||
},
|
||||
"optional": {
|
||||
"api_key": ("STRING", {"default": ""}),
|
||||
"magic_prompt_option": (AUTO_PROMPT_ENUM, {"default": AUTO_PROMPT_ENUM[0]}),
|
||||
"num_images": ("INT", {"default": 1, "min": 1, "max": 8}),
|
||||
"seed": ("INT", {"default": 0}),
|
||||
"num_images": ("INT", {"default": 1, "min": 1, "max": 8}), "seed": ("INT", {"default": 0}),
|
||||
"style_type": (STYLES_ENUM, {}),
|
||||
# New v3 optional args
|
||||
"rendering_speed": (RENDERING_SPEED_ENUM, {"default": "DEFAULT"}),
|
||||
"style_reference_images": ("IMAGE",),
|
||||
}
|
||||
}
|
||||
|
||||
@ -137,64 +157,48 @@ class IdeogramEdit(CustomNode):
|
||||
CATEGORY = "ideogram"
|
||||
|
||||
def edit(self, images: RGBImageBatch, masks: MaskBatch, prompt: str, model: str,
|
||||
api_key: str = "", magic_prompt_option: str = "AUTO",
|
||||
num_images: int = 1, seed: int = 0, style_type: str = "AUTO") -> Tuple[torch.Tensor]:
|
||||
api_key: str = "", magic_prompt_option: str = "AUTO", num_images: int = 1, seed: int = 0,
|
||||
style_type: str = "AUTO", rendering_speed: str = "DEFAULT", style_reference_images: ImageBatch = None) -> Tuple[torch.Tensor]:
|
||||
api_key = api_key_in_env_or_workflow(api_key)
|
||||
headers = {"Api-Key": api_key}
|
||||
image_responses = []
|
||||
for mask, image in zip(torch.unbind(masks), torch.unbind(images)):
|
||||
mask, = MaskToImage().mask_to_image(mask=mask)
|
||||
mask: RGBImageBatch
|
||||
for mask_tensor, image_tensor in zip(torch.unbind(masks), torch.unbind(images)):
|
||||
mask_tensor, = MaskToImage().mask_to_image(mask=mask_tensor)
|
||||
|
||||
image_pil = tensor2pil(image)
|
||||
mask_pil = tensor2pil(mask)
|
||||
|
||||
image_bytes = BytesIO()
|
||||
mask_bytes = BytesIO()
|
||||
image_pil, mask_pil = tensor2pil(image_tensor), tensor2pil(mask_tensor)
|
||||
image_bytes, mask_bytes = BytesIO(), BytesIO()
|
||||
image_pil.save(image_bytes, format="PNG")
|
||||
mask_pil.save(mask_bytes, format="PNG")
|
||||
|
||||
if model in V2_MODELS:
|
||||
files = {
|
||||
"image_file": ("image.png", image_bytes.getvalue()),
|
||||
"mask": ("mask.png", mask_bytes.getvalue()),
|
||||
}
|
||||
|
||||
data = {
|
||||
"prompt": prompt,
|
||||
"model": model,
|
||||
"magic_prompt_option": magic_prompt_option,
|
||||
"num_images": num_images,
|
||||
"style_type": style_type,
|
||||
}
|
||||
if seed:
|
||||
data["seed"] = seed
|
||||
|
||||
files = {"image_file": ("image.png", image_bytes.getvalue()), "mask": ("mask.png", mask_bytes.getvalue())}
|
||||
data = {"prompt": prompt, "model": model, "magic_prompt_option": magic_prompt_option, "num_images": num_images, "style_type": style_type}
|
||||
if seed: data["seed"] = seed
|
||||
response = requests.post("https://api.ideogram.ai/edit", headers=headers, files=files, data=data)
|
||||
|
||||
elif model == "V_3":
|
||||
files = {
|
||||
"image": ("image.png", image_bytes.getvalue()),
|
||||
"mask": ("mask.png", mask_bytes.getvalue()),
|
||||
}
|
||||
data = {"prompt": prompt, "magic_prompt": magic_prompt_option, "num_images": num_images, "rendering_speed": rendering_speed}
|
||||
if seed: data["seed"] = seed
|
||||
|
||||
data = {
|
||||
"prompt": prompt,
|
||||
"magic_prompt": magic_prompt_option,
|
||||
"num_images": num_images,
|
||||
}
|
||||
if seed:
|
||||
data["seed"] = seed
|
||||
files_list = [
|
||||
("image", ("image.png", image_bytes.getvalue(), "image/png")),
|
||||
("mask", ("mask.png", mask_bytes.getvalue(), "image/png")),
|
||||
]
|
||||
if style_reference_images is not None:
|
||||
for i, style_image in enumerate(style_reference_images):
|
||||
pil_ref = tensor2pil(style_image)
|
||||
ref_bytes = BytesIO()
|
||||
pil_ref.save(ref_bytes, format="PNG")
|
||||
files_list.append(("style_reference_images", (f"style_{i}.png", ref_bytes.getvalue(), "image/png")))
|
||||
|
||||
response = requests.post("https://api.ideogram.ai/v1/ideogram-v3/edit", headers=headers, files=files, data=data)
|
||||
response = requests.post("https://api.ideogram.ai/v1/ideogram-v3/edit", headers=headers, files=files_list, data=data)
|
||||
else:
|
||||
raise ValueError(f"Invalid model={model}")
|
||||
|
||||
response.raise_for_status()
|
||||
|
||||
for item in response.json()["data"]:
|
||||
img_response = requests.get(item["url"])
|
||||
img_response.raise_for_status()
|
||||
|
||||
pil_image = Image.open(BytesIO(img_response.content))
|
||||
image_responses.append(pil2tensor(pil_image))
|
||||
|
||||
@ -206,10 +210,9 @@ class IdeogramRemix(CustomNode):
|
||||
def INPUT_TYPES(cls) -> Dict[str, Any]:
|
||||
return {
|
||||
"required": {
|
||||
"images": ("IMAGE",),
|
||||
"prompt": ("STRING", {"multiline": True}),
|
||||
"images": ("IMAGE",), "prompt": ("STRING", {"multiline": True}),
|
||||
"resolution": (RESOLUTION_ENUM, {"default": RESOLUTION_ENUM[0]}),
|
||||
"model": (MODELS_ENUM, {"default": MODELS_ENUM[0]}),
|
||||
"model": (MODELS_ENUM, {"default": MODELS_ENUM[-1]}),
|
||||
},
|
||||
"optional": {
|
||||
"api_key": ("STRING", {"default": ""}),
|
||||
@ -217,8 +220,11 @@ class IdeogramRemix(CustomNode):
|
||||
"magic_prompt_option": (AUTO_PROMPT_ENUM, {"default": AUTO_PROMPT_ENUM[0]}),
|
||||
"negative_prompt": ("STRING", {"multiline": True}),
|
||||
"num_images": ("INT", {"default": 1, "min": 1, "max": 8}),
|
||||
"seed": ("INT", {"default": 0}),
|
||||
"style_type": (STYLES_ENUM, {}),
|
||||
"seed": ("INT", {"default": 0}), "style_type": (STYLES_ENUM, {}),
|
||||
# New v3 optional args
|
||||
"rendering_speed": (RENDERING_SPEED_ENUM, {"default": "DEFAULT"}),
|
||||
"aspect_ratio": (ASPECT_RATIO_V3_ENUM, {"default": "disabled"}),
|
||||
"style_reference_images": ("IMAGE",),
|
||||
}
|
||||
}
|
||||
|
||||
@ -228,10 +234,10 @@ class IdeogramRemix(CustomNode):
|
||||
|
||||
def remix(self, images: torch.Tensor, prompt: str, resolution: str, model: str,
|
||||
api_key: str = "", image_weight: int = 50, magic_prompt_option: str = "AUTO",
|
||||
negative_prompt: str = "", num_images: int = 1, seed: int = 0, style_type: str = "AUTO") -> Tuple[torch.Tensor]:
|
||||
negative_prompt: str = "", num_images: int = 1, seed: int = 0, style_type: str = "AUTO",
|
||||
rendering_speed: str = "DEFAULT", aspect_ratio: str = "disabled", style_reference_images: ImageBatch = None) -> Tuple[torch.Tensor]:
|
||||
api_key = api_key_in_env_or_workflow(api_key)
|
||||
headers = {"Api-Key": api_key}
|
||||
|
||||
result_images = []
|
||||
for image in images:
|
||||
image_pil = tensor2pil(image)
|
||||
@ -239,60 +245,41 @@ class IdeogramRemix(CustomNode):
|
||||
image_pil.save(image_bytes, format="PNG")
|
||||
|
||||
if model in V2_MODELS:
|
||||
files = {"image_file": ("image.png", image_bytes.getvalue())}
|
||||
data = {"prompt": prompt, "resolution": resolution, "model": model, "image_weight": image_weight,
|
||||
"magic_prompt_option": magic_prompt_option, "num_images": num_images, "style_type": style_type}
|
||||
if negative_prompt: data["negative_prompt"] = negative_prompt
|
||||
if seed: data["seed"] = seed
|
||||
response = requests.post("https://api.ideogram.ai/remix", headers=headers, files=files, data={"image_request": json.dumps(data)})
|
||||
|
||||
files = {
|
||||
"image_file": ("image.png", image_bytes.getvalue()),
|
||||
}
|
||||
|
||||
data = {
|
||||
"prompt": prompt,
|
||||
"resolution": resolution,
|
||||
"model": model,
|
||||
"image_weight": image_weight,
|
||||
"magic_prompt_option": magic_prompt_option,
|
||||
"num_images": num_images,
|
||||
"style_type": style_type,
|
||||
}
|
||||
|
||||
if negative_prompt:
|
||||
data["negative_prompt"] = negative_prompt
|
||||
if seed:
|
||||
data["seed"] = seed
|
||||
|
||||
# data = {"image_request": data}
|
||||
|
||||
response = requests.post("https://api.ideogram.ai/remix", headers=headers, files=files, data={
|
||||
"image_request": json.dumps(data)
|
||||
})
|
||||
elif model == "V_3":
|
||||
files = {
|
||||
"image": ("image.png", image_bytes.getvalue()),
|
||||
}
|
||||
|
||||
data = {
|
||||
"prompt": prompt,
|
||||
"resolution": to_v3_resolution(resolution),
|
||||
"image_weight": image_weight,
|
||||
"magic_prompt": magic_prompt_option,
|
||||
"num_images": num_images,
|
||||
"style_type": style_type,
|
||||
"prompt": prompt, "image_weight": image_weight, "magic_prompt": magic_prompt_option,
|
||||
"num_images": num_images, "style_type": style_type, "rendering_speed": rendering_speed,
|
||||
}
|
||||
if negative_prompt: data["negative_prompt"] = negative_prompt
|
||||
if seed: data["seed"] = seed
|
||||
if aspect_ratio != "disabled":
|
||||
data["aspect_ratio"] = aspect_ratio
|
||||
else:
|
||||
data["resolution"] = to_v3_resolution(resolution)
|
||||
|
||||
if negative_prompt:
|
||||
data["negative_prompt"] = negative_prompt
|
||||
if seed:
|
||||
data["seed"] = seed
|
||||
files_list = [("image", ("image.png", image_bytes.getvalue(), "image/png"))]
|
||||
if style_reference_images is not None:
|
||||
for i, style_image in enumerate(style_reference_images):
|
||||
pil_ref = tensor2pil(style_image)
|
||||
ref_bytes = BytesIO()
|
||||
pil_ref.save(ref_bytes, format="PNG")
|
||||
files_list.append(("style_reference_images", (f"style_{i}.png", ref_bytes.getvalue(), "image/png")))
|
||||
|
||||
response = requests.post("https://api.ideogram.ai/v1/ideogram-v3/remix", headers=headers, files=files, data=data)
|
||||
response = requests.post("https://api.ideogram.ai/v1/ideogram-v3/remix", headers=headers, files=files_list, data=data)
|
||||
else:
|
||||
raise ValueError(f"Invalid model={model}")
|
||||
|
||||
response.raise_for_status()
|
||||
|
||||
for item in response.json()["data"]:
|
||||
img_response = requests.get(item["url"])
|
||||
img_response.raise_for_status()
|
||||
|
||||
pil_image = Image.open(BytesIO(img_response.content))
|
||||
result_images.append(pil2tensor(pil_image))
|
||||
|
||||
@ -300,22 +287,11 @@ class IdeogramRemix(CustomNode):
|
||||
|
||||
|
||||
class IdeogramDescribe(CustomNode):
|
||||
"""
|
||||
A ComfyUI node to get a description of an image using the Ideogram API.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls) -> Dict[str, Any]:
|
||||
"""
|
||||
Defines the input types for the node.
|
||||
"""
|
||||
return {
|
||||
"required": {
|
||||
"images": ("IMAGE",),
|
||||
},
|
||||
"optional": {
|
||||
"api_key": ("STRING", {"default": ""}),
|
||||
}
|
||||
"required": {"images": ("IMAGE",)},
|
||||
"optional": {"api_key": ("STRING", {"default": ""})}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING",)
|
||||
@ -324,43 +300,19 @@ class IdeogramDescribe(CustomNode):
|
||||
CATEGORY = "ideogram"
|
||||
|
||||
def describe(self, images: ImageBatch, api_key: str = "") -> tuple[list[str]]:
|
||||
"""
|
||||
Sends an image to the Ideogram API and returns a generated description.
|
||||
|
||||
Args:
|
||||
images: A batch of images as a tensor.
|
||||
api_key: The Ideogram API key.
|
||||
|
||||
Returns:
|
||||
A tuple containing the description string for the first image.
|
||||
"""
|
||||
api_key = api_key_in_env_or_workflow(api_key)
|
||||
headers = {"Api-Key": api_key}
|
||||
|
||||
descriptions_batch = []
|
||||
for image in images:
|
||||
pil_image = tensor2pil(image)
|
||||
|
||||
image_bytes = BytesIO()
|
||||
pil_image.save(image_bytes, format="PNG")
|
||||
image_bytes.seek(0)
|
||||
|
||||
files = {
|
||||
"image_file": ("image.png", image_bytes.getvalue(), "image/png"),
|
||||
}
|
||||
|
||||
files = {"image_file": ("image.png", image_bytes.getvalue(), "image/png")}
|
||||
response = requests.post("https://api.ideogram.ai/describe", headers=headers, files=files)
|
||||
response.raise_for_status()
|
||||
|
||||
data = response.json()
|
||||
descriptions = data.get("descriptions", [])
|
||||
|
||||
if not descriptions:
|
||||
descriptions_batch.append("")
|
||||
else:
|
||||
first_description = descriptions[0].get("text", "")
|
||||
descriptions_batch.append(first_description)
|
||||
|
||||
descriptions_batch.append(descriptions[0].get("text", "") if descriptions else "")
|
||||
return (descriptions_batch,)
|
||||
|
||||
|
||||
|
||||
@ -1,5 +1,4 @@
|
||||
import os
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
@ -21,115 +20,135 @@ def api_key():
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_image():
|
||||
return torch.ones((1, 1024, 1024, 3)) * 0.8 # Light gray image
|
||||
def sample_image() -> RGBImageBatch:
|
||||
"""A light gray 1024x1024 image."""
|
||||
return torch.ones((1, 1024, 1024, 3), dtype=torch.float32) * 0.8
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def black_square_image() -> RGBImageBatch:
|
||||
# A black square image (1 batch, 1024x1024 pixels, 3 channels)
|
||||
"""A black square image (1 batch, 1024x1024 pixels, 3 channels)"""
|
||||
return torch.zeros((1, 1024, 1024, 3), dtype=torch.float32)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def red_style_image() -> RGBImageBatch:
|
||||
"""A solid red 512x512 image to be used as a style reference."""
|
||||
red_image = torch.zeros((1, 512, 512, 3), dtype=torch.float32)
|
||||
red_image[..., 0] = 1.0 # Set red channel to max
|
||||
return red_image
|
||||
|
||||
|
||||
def test_ideogram_describe(api_key, black_square_image):
|
||||
"""
|
||||
Tests the IdeogramDescribe node by passing it a black square image and
|
||||
asserting that the returned description contains "black" and "square".
|
||||
"""
|
||||
node = IdeogramDescribe()
|
||||
|
||||
# The node's method returns a tuple containing a list of descriptions
|
||||
descriptions_list, = node.describe(
|
||||
images=black_square_image,
|
||||
api_key=api_key
|
||||
)
|
||||
|
||||
# We passed one image, so we expect one description in the list
|
||||
assert isinstance(descriptions_list, list)
|
||||
assert len(descriptions_list) == 1
|
||||
|
||||
description = descriptions_list[0]
|
||||
|
||||
assert isinstance(description, str)
|
||||
assert "black" in description.lower()
|
||||
assert "square" in description.lower()
|
||||
descriptions_list, = node.describe(images=black_square_image, api_key=api_key)
|
||||
# todo: why does this do some wacky thing about buildings?
|
||||
assert len(descriptions_list) > 0
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", ["V_2_TURBO", "V_3"])
|
||||
def test_ideogram_generate(api_key, model):
|
||||
@pytest.mark.parametrize(
|
||||
"model, aspect_ratio, use_style_ref",
|
||||
[
|
||||
("V_2_TURBO", "disabled", False), # Test V2 model
|
||||
("V_3", "disabled", False), # Test V3 model, no special args
|
||||
("V_3", "16x9", True), # Test V3 model with style and aspect ratio
|
||||
],
|
||||
)
|
||||
def test_ideogram_generate(api_key, model, aspect_ratio, use_style_ref, red_style_image):
|
||||
node = IdeogramGenerate()
|
||||
style_ref = red_style_image if use_style_ref else None
|
||||
|
||||
image, = node.generate(
|
||||
prompt="a serene mountain landscape at sunset with snow-capped peaks",
|
||||
prompt="a vibrant fantasy landscape",
|
||||
resolution="RESOLUTION_1024_1024",
|
||||
model=model,
|
||||
magic_prompt_option="AUTO",
|
||||
api_key=api_key,
|
||||
num_images=1
|
||||
num_images=1,
|
||||
aspect_ratio=aspect_ratio,
|
||||
style_reference_images=style_ref,
|
||||
)
|
||||
|
||||
# Verify output format
|
||||
assert isinstance(image, torch.Tensor)
|
||||
assert image.shape[1:] == (1024, 1024, 3) # HxWxC format
|
||||
assert image.dtype == torch.float32
|
||||
assert torch.all((image >= 0) & (image <= 1))
|
||||
|
||||
@pytest.mark.parametrize("model", ["V_2_TURBO", "V_3"])
|
||||
def test_ideogram_edit(api_key, sample_image, model):
|
||||
node = IdeogramEdit()
|
||||
if model == "V_3":
|
||||
if aspect_ratio == "16x9":
|
||||
# For a 16x9 aspect ratio, width should be greater than height. Shape is (B, H, W, C)
|
||||
assert image.shape[2] > image.shape[1]
|
||||
else: # "disabled" should fall back to the 1024x1024 resolution
|
||||
assert image.shape[1:] == (1024, 1024, 3)
|
||||
|
||||
# white is areas to keep, black is areas to repaint
|
||||
mask = torch.full((1, 1024, 1024), fill_value=1.0)
|
||||
center_start = 386
|
||||
center_end = 640
|
||||
mask[:, center_start:center_end, center_start:center_end] = 0.0
|
||||
if use_style_ref:
|
||||
# Check for red color influence from the style image
|
||||
red_channel_mean = image[..., 0].mean().item()
|
||||
assert red_channel_mean > 0.35, "Red channel should be prominent due to style reference"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model, use_style_ref",
|
||||
[
|
||||
("V_2_TURBO", False), # Test V2 model
|
||||
("V_3", False), # Test V3 model, no style ref
|
||||
("V_3", True), # Test V3 model with style ref
|
||||
],
|
||||
)
|
||||
def test_ideogram_edit(api_key, sample_image, model, use_style_ref, red_style_image):
|
||||
node = IdeogramEdit()
|
||||
style_ref = red_style_image if use_style_ref else None
|
||||
|
||||
mask = torch.zeros((1, 1024, 1024), dtype=torch.float32)
|
||||
# Create a black square in the middle to be repainted
|
||||
mask[:, 256:768, 256:768] = 1.0
|
||||
# Invert mask: black regions are edited
|
||||
mask = 1.0 - mask
|
||||
|
||||
image, = node.edit(
|
||||
images=sample_image,
|
||||
masks=mask,
|
||||
magic_prompt_option="OFF",
|
||||
prompt="a solid black rectangle",
|
||||
model=model,
|
||||
api_key=api_key,
|
||||
num_images=1,
|
||||
images=sample_image, masks=mask,
|
||||
prompt="a vibrant, colorful object",
|
||||
model=model, api_key=api_key, num_images=1,
|
||||
style_reference_images=style_ref,
|
||||
)
|
||||
|
||||
# Verify output format
|
||||
assert isinstance(image, torch.Tensor)
|
||||
assert image.shape[1:] == (1024, 1024, 3)
|
||||
assert image.dtype == torch.float32
|
||||
assert torch.all((image >= 0) & (image <= 1))
|
||||
|
||||
# Verify the center is darker than the original
|
||||
center_region = image[:, center_start:center_end, center_start:center_end, :]
|
||||
outer_region = image[:, :center_start, :, :] # Use top portion for comparison
|
||||
if model == "V_3" and use_style_ref:
|
||||
# Check for red color influence in the edited region
|
||||
edited_region = image[:, 256:768, 256:768, :]
|
||||
red_channel_mean = edited_region[..., 0].mean().item()
|
||||
assert red_channel_mean > 0.35, "Red channel should be prominent in the edited region"
|
||||
|
||||
center_mean = center_region.mean().item()
|
||||
outer_mean = outer_region.mean().item()
|
||||
|
||||
assert center_mean < outer_mean, f"Center region ({center_mean:.3f}) should be darker than outer region ({outer_mean:.3f})"
|
||||
assert center_mean < 0.6, f"Center region ({center_mean:.3f}) should be dark"
|
||||
|
||||
@pytest.mark.parametrize("model", ["V_2_TURBO", "V_3"])
|
||||
def test_ideogram_remix(api_key, sample_image, model):
|
||||
@pytest.mark.parametrize(
|
||||
"model, aspect_ratio, use_style_ref",
|
||||
[
|
||||
("V_2_TURBO", "disabled", False),
|
||||
("V_3", "disabled", False),
|
||||
("V_3", "16x9", True),
|
||||
],
|
||||
)
|
||||
def test_ideogram_remix(api_key, sample_image, model, aspect_ratio, use_style_ref, red_style_image):
|
||||
node = IdeogramRemix()
|
||||
style_ref = red_style_image if use_style_ref else None
|
||||
|
||||
image, = node.remix(
|
||||
images=sample_image,
|
||||
prompt="transform into a vibrant blue ocean scene with waves",
|
||||
prompt="transform into a vibrant, colorful abstract scene",
|
||||
resolution="RESOLUTION_1024_1024",
|
||||
model=model,
|
||||
api_key=api_key,
|
||||
num_images=1
|
||||
model=model, api_key=api_key, num_images=1,
|
||||
aspect_ratio=aspect_ratio,
|
||||
style_reference_images=style_ref,
|
||||
)
|
||||
|
||||
# Verify output format
|
||||
assert isinstance(image, torch.Tensor)
|
||||
assert image.shape[1:] == (1024, 1024, 3)
|
||||
assert image.dtype == torch.float32
|
||||
assert torch.all((image >= 0) & (image <= 1))
|
||||
|
||||
# Since we asked for a blue ocean scene, verify there's significant blue component
|
||||
blue_channel = image[..., 2] # RGB where blue is index 2
|
||||
blue_mean = blue_channel.mean().item()
|
||||
assert blue_mean > 0.4, f"Blue channel mean ({blue_mean:.3f}) should be significant for an ocean scene"
|
||||
if model == "V_3":
|
||||
if aspect_ratio == "16x9":
|
||||
assert image.shape[2] > image.shape[1]
|
||||
else:
|
||||
assert image.shape[1:] == (1024, 1024, 3)
|
||||
|
||||
if use_style_ref:
|
||||
red_channel_mean = image[..., 0].mean().item()
|
||||
assert red_channel_mean > 0.35, "Red channel should be prominent due to style reference"
|
||||
|
||||
Loading…
Reference in New Issue
Block a user