pantaleone-ai-Comfy-Firefly/firefly_node.py

import requests
import time
import torch
import numpy as np
from PIL import Image
from io import BytesIO

class AdobeFireflyTextToImage:
    """
    A custom node for ComfyUI to generate images using Adobe Firefly Services API (v3).
    Handles authentication (OAuth Server-to-Server), job submission, polling, and image retrieval.
    """

    def __init__(self):
        self.base_url = "https://firefly-api.adobe.io"
        self.auth_url = "https://ims-na1.adobelogin.com/ims/token/v3"
        self.cached_token = None
        self.token_expiry = 0

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "client_id": ("STRING", {"multiline": False, "default": "", "placeholder": "Enter Adobe Client ID (API Key)"}),
                "client_secret": ("STRING", {"multiline": False, "default": "", "placeholder": "Enter Adobe Client Secret"}),
                "prompt": ("STRING", {"multiline": True, "default": "A futuristic city with flying cars, cinematic lighting", "placeholder": "Describe your image..."}),
                "aspect_ratio": (["square", "landscape", "portrait", "widescreen"], {"default": "square"}),
                "content_class": (["photo", "art", "graphic"], {"default": "art"}),
                "visual_intensity": ("INT", {"default": 6, "min": 1, "max": 10}),
                "style_preset": (
                    [
                        "none", "photorealistic", "digital_art", "oil_painting", "neon", 
                        "cinematic", "3d_model", "origami", "pencil_sketch", "pixel_art"
                    ], 
                    {"default": "none"}
                ),
            },
            "optional": {
                "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
            }
        }

    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("image",)
    FUNCTION = "generate_image"
    CATEGORY = "Adobe Firefly"

    def get_access_token(self, client_id, client_secret):
        """
        Exchanges Client ID and Secret for a bearer token.
        Caches the token until it is close to expiry.
        """
        current_time = time.time()
        
        # Return cached token if still valid (buffer of 60 seconds)
        if self.cached_token and current_time < (self.token_expiry - 60):
            return self.cached_token

        headers = {
            "Content-Type": "application/x-www-form-urlencoded"
        }
        
        # Standard scope for Firefly Services
        data = {
            "grant_type": "client_credentials",
            "client_id": client_id,
            "client_secret": client_secret,
            "scope": "openid,AdobeID,firefly_api,ff_apis"
        }

        try:
            print(f"Requesting new Access Token from {self.auth_url}...")
            response = requests.post(self.auth_url, headers=headers, data=data)
            response.raise_for_status()
            
            json_response = response.json()
            self.cached_token = json_response["access_token"]
            # 'expires_in' is usually in milliseconds or seconds
            expires_in = int(json_response.get("expires_in", 3600)) 
            self.token_expiry = current_time + expires_in
            
            print("Successfully retrieved Adobe Access Token.")
            return self.cached_token
            
        except requests.exceptions.RequestException as e:
            error_msg = f"Failed to authenticate with Adobe: {str(e)}"
            if e.response is not None:
                error_msg += f"\nResponse: {e.response.text}"
            raise Exception(error_msg)

    def generate_image(self, client_id, client_secret, prompt, aspect_ratio, content_class, visual_intensity, style_preset, seed):
        
        if not client_id or not client_secret:
            raise ValueError("Adobe Client ID and Client Secret are required.")

        # 1. Authenticate
        access_token = self.get_access_token(client_id, client_secret)

        # 2. Prepare Payload
        
        payload = {
            "prompt": prompt,
            "n": 1, 
            "size": {
                "width": 2048,
                "height": 2048
            },
            "photoSettings": {
                "aperture": 5.6,
                "shutterSpeed": 0.005,
                "iso": 400
            },
             "contentClass": content_class,
             "visualIntensity": visual_intensity,
             "seeds": [seed]
        }

        # Adjust size based on aspect ratio request
        if aspect_ratio == "square":
            payload["size"] = {"width": 2048, "height": 2048}
        elif aspect_ratio == "landscape":
            payload["size"] = {"width": 2688, "height": 1536}
        elif aspect_ratio == "portrait":
            payload["size"] = {"width": 1536, "height": 2688}
        elif aspect_ratio == "widescreen":
            payload["size"] = {"width": 2304, "height": 1296}

        if style_preset != "none":
            payload["styles"] = {
                "presets": [style_preset]
            }

        headers = {
            "X-Api-Key": client_id,
            "Authorization": f"Bearer {access_token}",
            "Content-Type": "application/json",
            "Accept": "application/json"
        }

        # 3. Submit Job
        generate_url = f"{self.base_url}/v3/images/generate"
        
        try:
            print(f"Submitting generation job to {generate_url}...")
            response = requests.post(generate_url, json=payload, headers=headers)
            
            if response.status_code != 200:
                raise Exception(f"Firefly API Error ({response.status_code}): {response.text}")

            response_data = response.json()
            
            image_url = None
            
            if "outputs" in response_data:
                image_url = response_data["outputs"][0]["image"]["url"]
            
            if not image_url:
                raise Exception(f"No image URL found in response. Response: {response_data}")

            # 4. Download Image
            print(f"Downloading image from {image_url}...")
            image_response = requests.get(image_url)
            image_response.raise_for_status()
            
            # 5. Convert to ComfyUI Format
            img = Image.open(BytesIO(image_response.content))
            img = img.convert("RGB")
            img_np = np.array(img).astype(np.float32) / 255.0
            img_tensor = torch.from_numpy(img_np)[None,]
            
            return (img_tensor,)

        except Exception as e:
            print(f"Error in Adobe Firefly Node: {e}")
            raise e