sumethweer-cmd-comfyui-gcs-uploader/gcs_upload_node.py

"""
GCS Upload Node for ComfyUI
============================
Uploads generated images directly from ComfyUI (RunPod) to Google Cloud Storage,
then notifies your Supabase database — completely bypassing the Supabase Edge Function.

This eliminates Supabase Egress costs for image data.

SETUP:
  1. Copy config.json.example → config.json in this directory.
  2. Fill in your Supabase URL and Service Role Key.
  3. Set your GCS credentials in the web dashboard (they are fetched at runtime).
"""

import os
import io
import json
import time
import traceback
import importlib.util
from pathlib import Path

import requests
import numpy as np
from PIL import Image

# ─── Config Loading ────────────────────────────────────────────────────────────

CONFIG_PATH = Path(__file__).parent / "config.json"

def load_config() -> dict:
    if not CONFIG_PATH.exists():
        raise FileNotFoundError(
            f"[GCS Uploader] config.json not found at {CONFIG_PATH}. "
            "Copy config.json.example → config.json and fill in your Supabase credentials."
        )
    with open(CONFIG_PATH, "r") as f:
        return json.load(f)


# ─── Supabase Helpers ──────────────────────────────────────────────────────────

def fetch_gcs_config(supabase_url: str, service_key: str) -> dict:
    """Fetches GCS credentials from the system_configs table in Supabase."""
    headers = {
        "apikey": service_key,
        "Authorization": f"Bearer {service_key}",
    }
    params = {
        "key_name": "in.(GCS_PROJECT_ID,GCS_CLIENT_EMAIL,GCS_PRIVATE_KEY,GCS_BUCKET_NAME)",
        "select": "key_name,key_value",
    }
    resp = requests.get(f"{supabase_url}/rest/v1/system_configs", headers=headers, params=params, timeout=10)
    resp.raise_for_status()

    rows = resp.json()
    config = {r["key_name"]: r["key_value"] for r in rows}

    required = ["GCS_PROJECT_ID", "GCS_CLIENT_EMAIL", "GCS_PRIVATE_KEY", "GCS_BUCKET_NAME"]
    missing = [k for k in required if not config.get(k)]
    if missing:
        raise ValueError(f"[GCS Uploader] Missing GCS config in system_configs: {missing}")

    # Fix escaped newlines in private key
    config["GCS_PRIVATE_KEY"] = config["GCS_PRIVATE_KEY"].replace("\\n", "\n")
    return config


def mark_job_completed(supabase_url: str, service_key: str, job_id: str, content_item_id: str,
                       image_type: str, slot_index: int, gcs_url: str, file_name: str):
    """Updates production_jobs and inserts into generated_images."""
    headers = {
        "apikey": service_key,
        "Authorization": f"Bearer {service_key}",
        "Content-Type": "application/json",
        "Prefer": "resolution=merge-duplicates",
    }

    # 1. Upsert into generated_images
    image_payload = {
        "content_item_id": content_item_id,
        "image_type": image_type,
        "file_path": gcs_url,
        "file_name": file_name,
        "slot_index": slot_index,
        "status": "Generated",
        "media_type": "image",
    }
    img_resp = requests.post(
        f"{supabase_url}/rest/v1/generated_images",
        headers={**headers, "Prefer": "resolution=merge-duplicates,return=minimal"},
        json=image_payload,
        timeout=10,
    )
    img_resp.raise_for_status()

    # 2. Mark production_job as Completed
    job_resp = requests.patch(
        f"{supabase_url}/rest/v1/production_jobs?id=eq.{job_id}",
        headers=headers,
        json={"status": "Completed", "updated_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())},
        timeout=10,
    )
    job_resp.raise_for_status()


def log_to_supabase(supabase_url: str, service_key: str, level: str, phase: str, message: str):
    """Writes a log entry to system_logs table."""
    try:
        headers = {
            "apikey": service_key,
            "Authorization": f"Bearer {service_key}",
            "Content-Type": "application/json",
        }
        requests.post(
            f"{supabase_url}/rest/v1/system_logs",
            headers=headers,
            json={"level": level, "phase": phase, "message": message},
            timeout=5,
        )
    except Exception:
        pass  # Logging failure should not fail the node


# ─── GCS Upload ───────────────────────────────────────────────────────────────

def upload_to_gcs(image_bytes: bytes, gcs_config: dict, object_path: str) -> str:
    """Uploads raw image bytes to GCS and returns the public URL."""
    from google.cloud import storage
    from google.oauth2 import service_account

    credentials = service_account.Credentials.from_service_account_info(
        {
            "type": "service_account",
            "project_id": gcs_config["GCS_PROJECT_ID"],
            "private_key_id": "key",
            "private_key": gcs_config["GCS_PRIVATE_KEY"],
            "client_email": gcs_config["GCS_CLIENT_EMAIL"],
            "auth_uri": "https://accounts.google.com/o/oauth2/auth",
            "token_uri": "https://oauth2.googleapis.com/token",
        },
        scopes=["https://www.googleapis.com/auth/devstorage.read_write"],
    )

    client = storage.Client(project=gcs_config["GCS_PROJECT_ID"], credentials=credentials)
    bucket = client.bucket(gcs_config["GCS_BUCKET_NAME"])
    blob = bucket.blob(object_path)
    blob.upload_from_string(image_bytes, content_type="image/png")

    return f"https://storage.googleapis.com/{gcs_config['GCS_BUCKET_NAME']}/{object_path}"


# ─── ComfyUI Node Definition ──────────────────────────────────────────────────

class GCSUploadNode:
    """
    ComfyUI node that uploads the generated image directly to Google Cloud Storage
    and notifies your Supabase database — no Supabase Egress cost!
    """

    CATEGORY = "AI Influencer / Storage"
    FUNCTION = "upload"
    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("gcs_url",)
    OUTPUT_NODE = True

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
                "job_id": ("STRING", {"default": "", "multiline": False}),
                "content_item_id": ("STRING", {"default": "", "multiline": False}),
                "image_type": (["SFW", "NSFW"],),
                "slot_index": ("INT", {"default": 0, "min": 0, "max": 16}),
            },
        }

    def upload(
        self,
        image,
        job_id: str,
        content_item_id: str,
        image_type: str,
        slot_index: int,
    ):
        print(f"\n[GCS Uploader] 🚀 Starting upload for job {job_id} (slot {slot_index}, {image_type})")

        try:
            # 1. Load config
            cfg = load_config()
            supabase_url = cfg["supabase_url"].rstrip("/")
            service_key = cfg["supabase_service_role_key"]

            # 2. Fetch GCS credentials from Supabase
            print("[GCS Uploader] 🔑 Fetching GCS config from Supabase...")
            gcs_config = fetch_gcs_config(supabase_url, service_key)

            # 3. Convert ComfyUI tensor → PNG bytes
            # image tensor shape: (batch, H, W, C) float32 in [0,1]
            img_np = (image[0].cpu().numpy() * 255).clip(0, 255).astype(np.uint8)
            pil_img = Image.fromarray(img_np, mode="RGB")
            buffer = io.BytesIO()
            pil_img.save(buffer, format="PNG")
            image_bytes = buffer.getvalue()
            print(f"[GCS Uploader] 🖼️  Image size: {len(image_bytes) / 1024:.1f} KB")

            # 4. Build GCS object path
            object_path = f"images/{content_item_id}/{image_type}/{job_id}_{slot_index}.png"

            # 5. Upload to GCS
            print(f"[GCS Uploader] 📤 Uploading to GCS: {object_path}")
            gcs_url = upload_to_gcs(image_bytes, gcs_config, object_path)
            print(f"[GCS Uploader] ✅ Uploaded: {gcs_url}")

            # 6. Update Supabase DB
            print("[GCS Uploader] 💾 Updating Supabase DB...")
            mark_job_completed(
                supabase_url, service_key,
                job_id=job_id,
                content_item_id=content_item_id,
                image_type=image_type,
                slot_index=slot_index,
                gcs_url=gcs_url,
                file_name=object_path,
            )
            print("[GCS Uploader] ✅ Job marked Completed in DB!")

            # 7. Log success
            log_to_supabase(
                supabase_url, service_key,
                "SUCCESS", "ComfyUI Node",
                f"Job {job_id} completed. Image uploaded to GCS: {object_path}"
            )

            return (gcs_url,)

        except Exception as e:
            err_msg = f"[GCS Uploader] ❌ Upload failed for job {job_id}: {e}"
            print(err_msg)
            traceback.print_exc()

            # Try to mark job as Failed in DB
            try:
                cfg = load_config()
                supabase_url = cfg["supabase_url"].rstrip("/")
                service_key = cfg["supabase_service_role_key"]
                headers = {
                    "apikey": service_key,
                    "Authorization": f"Bearer {service_key}",
                    "Content-Type": "application/json",
                }
                requests.patch(
                    f"{supabase_url}/rest/v1/production_jobs?id=eq.{job_id}",
                    headers=headers,
                    json={
                        "status": "Failed",
                        "error_message": str(e),
                        "updated_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
                    },
                    timeout=5,
                )
                log_to_supabase(supabase_url, service_key, "ERROR", "ComfyUI Node", err_msg)
            except Exception:
                pass

            raise RuntimeError(err_msg) from e