diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..411932f --- /dev/null +++ b/__init__.py @@ -0,0 +1,11 @@ +from .rd_api_node import RetroDiffusionAPINode + +NODE_CLASS_MAPPINGS = { + "Retro Diffusion API Node": RetroDiffusionAPINode +} + +NODE_DISPLAY_NAMES_MAPPINGS = { + "Retro Diffusion API Node": "Retro Diffusion API Node" +} + +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] diff --git a/rd_api_node.py b/rd_api_node.py new file mode 100644 index 0000000..bec8a10 --- /dev/null +++ b/rd_api_node.py @@ -0,0 +1,130 @@ +from PIL import Image +import numpy as np +import torch +import base64 +from io import BytesIO +import requests +import os +from dotenv import load_dotenv + +load_dotenv() + +def base64_to_image(base64_string): + # Convert to PIL image + image_data = base64.b64decode(base64_string) + image = Image.open(BytesIO(image_data)) + image = image.convert("RGB") + + # Convert to tensor format + image_tensor = torch.from_numpy(np.array(image)).float() / 255.0 + image_tensor = image_tensor.unsqueeze(0) + + return image_tensor + +def image_to_base64(image_input): + # Convert from tensor format + image_np = image_input.squeeze(0).numpy() * 255.0 + image_np = image_np.clip(0, 255).astype(np.uint8) + + # Convert to PIL image + image = Image.fromarray(image_np, 'RGB') + buffer = BytesIO() + image.save(buffer, format="PNG") + + image_bytes = buffer.getvalue() + + base64_string = base64.b64encode(image_bytes).decode('utf-8') + + return base64_string + +class RetroDiffusionAPINode: + @classmethod + def INPUT_TYPES(cls): + prompt_styles = ["rd_fast__default", "rd_fast__retro", "rd_fast__simple","rd_fast__detailed", + "rd_fast__anime", "rd_fast__game_asset", "rd_fast__portrait", "rd_fast__texture", + "rd_fast__ui", "rd_fast__item_sheet", "rd_fast__character_turnaround", "rd_fast__1_bit", + "rd_fast__low_res", "rd_fast__mc_item", "rd_fast__mc_texture", "rd_fast__no_style", + "rd_plus__default", "rd_plus__retro", "rd_plus__watercolor", "rd_plus__textured", + "rd_plus__cartoon", "rd_plus__ui_element", "rd_plus__item_sheet", "rd_plus__character_turnaround", + "rd_plus__topdown_map", "rd_plus__topdown_asset", "rd_plus__isometric", "rd_plus__isometric_asset", + "rd_plus__classic", "rd_plus__low_res", "rd_plus__mc_item", "rd_plus__mc_texture", + "animation__four_angle_walking", "animation__walking_and_idle", "animation__vfx"] + return { + "required": { + #"api_key": ("STRING", {"default": "YOUR_API_KEY", "multiline": False}), + "seed": ("INT", {"default": 0, "min": 0, "max": 2147483647, "step": 1}), + "width": ("INT", {"default": 128, "min": 16, "max": 512, "step": 8}), + "height": ("INT", {"default": 128, "min": 16, "max": 512, "step": 8}), + "prompt_style": (prompt_styles,), + "remove_bg": ("BOOLEAN", {"default": False}), + "tile_x": ("BOOLEAN", {"default": False}), + "tile_y": ("BOOLEAN", {"default": False}), + "return_spritesheet_for_animations": ("BOOLEAN", {"default": True}), + #"num_images": ("INT", {"default": 1, "min": 1, "max": 16, "step": 1}), + "prompt": ("STRING", {"default": "", "multiline": True}), + "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}) + }, + "optional": { + "input_image": ("IMAGE",) + } + } + + RETURN_TYPES = ("IMAGE", "INT", "INT") + RETURN_NAMES = ("image", "credit_cost", "remaining_credits") + FUNCTION = "generate_image" + CATEGORY = "RetroDiffusion" + + def generate_image(self, seed, width, height, prompt_style, remove_bg, tile_x, tile_y, return_spritesheet_for_animations, prompt, denoise, input_image=None): + url = "https://api.retrodiffusion.ai/v1/inferences" + + headers = { + "X-RD-Token": os.getenv("RD_API_KEY"), + } + + payload = { + "seed": seed, + "width": width, + "height": height, + "prompt_style": prompt_style, + "remove_bg": remove_bg, + "tile_x": tile_x, + "tile_y": tile_y, + "return_spritesheet": return_spritesheet_for_animations, + "prompt": prompt, + "num_images": 1, + } + + if input_image is not None: + payload["input_image"] = image_to_base64(input_image) + payload["strength"] = denoise + + try: + response = requests.post(url, headers=headers, json=payload) + response.raise_for_status() + data = response.json() + + # Extract response info + credit_cost = data.get("credit_cost", 0) + print(f"Credit cost: {credit_cost}") + remaining_credits = data.get("remaining_credits", 0) + print(f"Remaining credits: {remaining_credits}") + + # Handle base64 image + base64_images = data.get("base64_images", []) + if not base64_images: + raise ValueError("No images returned from RetroDiffusion API") + + # Convert base64 to image + base64_image = base64_images[0] # Take first image + image_tensor = base64_to_image(base64_image) + + return (image_tensor, credit_cost, remaining_credits) + + except requests.exceptions.RequestException as e: + print(f"API request failed: {e}") + # Return empty tensor and 0 cost on error + return (torch.zeros(1, 64, 64, 3), 0, -1) + except Exception as e: + print(f"Error processing response: {e}") + # Return empty tensor and 0 cost on error + return (torch.zeros(1, 64, 64, 3), 0, -1) \ No newline at end of file