deerflow2/skills/public/image-generation/scripts/generate.py

import base64
import json
import os
import time
from typing import List

import requests
from PIL import Image
from dotenv import load_dotenv

load_dotenv()


def validate_image(image_path: str) -> bool:
    """
    Validate if an image file can be opened and is not corrupted.
    
    Args:
        image_path: Path to the image file
        
    Returns:
        True if the image is valid and can be opened, False otherwise
    """
    try:
        with Image.open(image_path) as img:
            img.verify()
        with Image.open(image_path) as img:
            img.load()
        return True
    except Exception as e:
        print(f"Warning: Image '{image_path}' is invalid or corrupted: {e}")
        return False


def submit_generation_task(prompt: str, aspect_ratio: str = "16:9", output_format: str = "png") -> str:
    """
    Submit image generation task to RunningHub API.
    
    Args:
        prompt: Text prompt for image generation
        aspect_ratio: Aspect ratio of the generated image
        output_format: Output image format (png, jpeg, webp)
        
    Returns:
        Task ID for tracking the generation
    """
    api_key = os.getenv("RUNNINGHUB_API_KEY")
    if not api_key:
        raise Exception("RUNNINGHUB_API_KEY environment variable is not set")
    
    url = "https://www.runninghub.cn/openapi/v2/rhart-image/z-image/turbo-lora"
    
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }
    
    payload = {
        "prompt": prompt,
        "aspectRatio": aspect_ratio,
        "lora_name": "Z-Image _ 清纯高颜值_脸模版V1.0.safetensors",
        "lora_strength": 1,
        "outputFormat": output_format
    }
    
    response = requests.post(url, headers=headers, json=payload)
    response.raise_for_status()
    result = response.json()
    
    if result.get("status") not in ["QUEUED", "RUNNING", "SUCCESS"]:
        raise Exception(f"Task submission failed: {result.get('errorMessage', 'Unknown error')}")
    
    return result.get("taskId")


def query_task_status(task_id: str) -> dict:
    """
    Query the status of a generation task.
    
    Args:
        task_id: Task ID to query
        
    Returns:
        Task status information
    """
    api_key = os.getenv("RUNNINGHUB_API_KEY")
    if not api_key:
        raise Exception("RUNNINGHUB_API_KEY environment variable is not set")
    
    url = "https://www.runninghub.cn/openapi/v2/query"
    
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }
    
    payload = {
        "taskId": task_id
    }
    
    response = requests.post(url, headers=headers, json=payload)
    response.raise_for_status()
    return response.json()


def download_image(url: str, output_path: str) -> None:
    """
    Download image from URL and save to file.
    
    Args:
        url: Image URL to download
        output_path: Local path to save the image
    """
    response = requests.get(url, stream=True)
    response.raise_for_status()
    
    with open(output_path, "wb") as f:
        for chunk in response.iter_content(chunk_size=8192):
            f.write(chunk)


def generate_image(
    prompt_file: str,
    reference_images: List[str],
    output_file: str,
    aspect_ratio: str = "16:9",
) -> str:
    """
    Generate image using RunningHub API.
    
    Args:
        prompt_file: Path to JSON prompt file
        reference_images: List of reference image paths (currently not supported by RunningHub API)
        output_file: Output path for generated image
        aspect_ratio: Aspect ratio of the generated image
        
    Returns:
        Success message with output file path
    """
    with open(prompt_file, "r", encoding="utf-8") as f:
        prompt_data = json.load(f)
    
    if reference_images:
        print("Note: RunningHub API does not support reference images in this version. Reference images will be ignored.")
    
    prompt_text = prompt_data.get("prompt", "")
    if not prompt_text:
        prompt_text = json.dumps(prompt_data, ensure_ascii=False)
    
    output_format = "png"
    if output_file.lower().endswith(".jpg") or output_file.lower().endswith(".jpeg"):
        output_format = "jpeg"
    elif output_file.lower().endswith(".webp"):
        output_format = "webp"
    
    print(f"Submitting generation task...")
    task_id = submit_generation_task(prompt_text, aspect_ratio, output_format)
    print(f"Task submitted successfully. Task ID: {task_id}")
    
    max_retries = 60
    retry_interval = 2
    
    for attempt in range(max_retries):
        print(f"Checking task status... (Attempt {attempt + 1}/{max_retries})")
        status_result = query_task_status(task_id)
        status = status_result.get("status")
        
        if status == "SUCCESS":
            print("Task completed successfully!")
            results = status_result.get("results", [])
            if results and len(results) > 0:
                image_url = results[0].get("url")
                if image_url:
                    print(f"Downloading image from: {image_url}")
                    download_image(image_url, output_file)
                    return f"Successfully generated image to {output_file}"
                else:
                    raise Exception("No image URL found in task results")
            else:
                raise Exception("No results found in task response")
        elif status == "FAILED":
            error_msg = status_result.get("errorMessage", "Unknown error")
            raise Exception(f"Task failed: {error_msg}")
        elif status in ["QUEUED", "RUNNING"]:
            print(f"Task status: {status}. Waiting...")
            time.sleep(retry_interval)
        else:
            raise Exception(f"Unknown task status: {status}")
    
    raise Exception(f"Task did not complete within {max_retries * retry_interval} seconds")


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Generate images using Gemini API")
    parser.add_argument(
        "--prompt-file",
        required=True,
        help="Absolute path to JSON prompt file",
    )
    parser.add_argument(
        "--reference-images",
        nargs="*",
        default=[],
        help="Absolute paths to reference images (space-separated)",
    )
    parser.add_argument(
        "--output-file",
        required=True,
        help="Output path for generated image",
    )
    parser.add_argument(
        "--aspect-ratio",
        required=False,
        default="16:9",
        help="Aspect ratio of the generated image",
    )

    args = parser.parse_args()

    try:
        print(
            generate_image(
                args.prompt_file,
                args.reference_images,
                args.output_file,
                args.aspect_ratio,
            )
        )
    except Exception as e:
        print(f"Error while generating image: {e}")