flux_kontext_lib/image_generator.py

import requests
import base64
import json
import os
from io import BytesIO
from typing import Optional, Dict

# Pillow is required for image format conversion and normalization.
# Please install it using: pip install Pillow
try:
    from PIL import Image
except ImportError:
    print("Pillow library not found. Please install it using: pip install Pillow")
    exit()

# --- Configuration ---
IMGBB_API_KEY = os.getenv("IMGBB_API_KEY")

def upload_to_imgbb(image_path: str, file_name: str) -> Optional[str]:
    """
    Uploads the image at image_path to ImgBB.
    Returns the public URL or None on failure.
    """
    if not IMGBB_API_KEY:
        print("Warning: IMGBB_API_KEY not set, skipping upload")
        return None

    print(f"Uploading {file_name} to ImgBB...")
    try:
        with open(image_path, 'rb') as f:
            files = {"image": (file_name, f.read())}
        resp = requests.post(
            "https://api.imgbb.com/1/upload",
            params={"key": IMGBB_API_KEY},
            files=files,
            timeout=20
        )
        resp.raise_for_status()
        data = resp.json().get("data", {})
        url = data.get("url")
        if url:
            print(f"Successfully uploaded to ImgBB: {url}")
            return url
        else:
            print(f"Error: ImgBB API response missing 'url'. Response: {resp.json()}")
            return None
    except requests.exceptions.RequestException as e:
        print(f"Error: ImgBB upload failed: {e}")
        return None

def _download_image_from_url(image_url: str, save_path: str) -> bool:
    """
    Downloads an image from a URL and saves it to a local path.

    Args:
        image_url: The URL of the image to download.
        save_path: The local path to save the downloaded image.

    Returns:
        True if the download was successful, False otherwise.
    """
    print(f"Downloading generated image from: {image_url}")
    try:
        image_response = requests.get(image_url, stream=True, timeout=30)
        # Check if the download request was successful.
        if image_response.status_code == 200:
            content_type = image_response.headers.get('Content-Type', '')
            if 'image' in content_type:
                with open(save_path, 'wb') as f:
                    for chunk in image_response.iter_content(1024):
                        f.write(chunk)
                print(f"Image successfully saved to {save_path}")
                return True
            else:
                print(f"Error: Content at URL is not an image. Content-Type: {content_type}")
                return False
        else:
            print(f"Error: Failed to download image. Status code: {image_response.status_code}")
            return False
    except requests.exceptions.RequestException as e:
        print(f"An error occurred during image download: {e}")
        return False

def generate_image(
    prompt_text: str,
    image_path: str,
    download_path: Optional[str] = None
) -> Optional[Dict]:
    """
    Sends a request to the image generation API, optionally downloads the result,
    and uploads it to ImgBB.

    Args:
        prompt_text: The instructional text for image modification.
        image_path: The file path to the input image (any common format).
        download_path: If provided, the path to save the generated image.

    Returns:
        A dictionary of the JSON response from the API, including the ImgBB URL,
        or None on error.
    """
    url = "https://kontext-chat.replicate.dev/generate-image"

    try:
        # --- Image Normalization Step ---
        with Image.open(image_path) as img:
            if img.mode != 'RGB':
                img = img.convert('RGB')
            with BytesIO() as output_buffer:
                img.save(
                    output_buffer,
                    format="JPEG",
                    quality=95,
                    subsampling=0,
                    progressive=False
                )
                image_bytes = output_buffer.getvalue()
    except FileNotFoundError:
        print(f"Error: Image file not found at {image_path}")
        return None
    except Exception as e:
        print(f"Error processing image file. Ensure it's a valid image. Details: {e}")
        return None

    encoded_string = base64.b64encode(image_bytes).decode('utf-8')
    input_image_data_uri = f"data:image/jpeg;base64,{encoded_string}"

    payload = {
        "prompt": prompt_text,
        "input_image": input_image_data_uri
    }

    headers = {
        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:139.0) Gecko/20100101 Firefox/139.0",
        "Accept": "*/*",
        "Content-Type": "application/json",
        "Referer": "https://kontext-chat.replicate.dev/",
        "Origin": "https://kontext-chat.replicate.dev",
    }

    try:
        response = requests.post(url, headers=headers, data=json.dumps(payload), timeout=45)
        response.raise_for_status()
        api_response_data = response.json()

        imgbb_url = None  # Initialize imgbb_url

        # --- Download & Upload Logic ---
        if download_path and isinstance(api_response_data, dict):
            image_url = api_response_data.get("imageUrl")
            if image_url:
                if _download_image_from_url(image_url, download_path):
                    # If download is successful, upload the saved file to ImgBB
                    file_name = os.path.basename(download_path)
                    imgbb_url = upload_to_imgbb(download_path, file_name)
            else:
                print("Warning: 'imageUrl' not found in API response, cannot download or upload image.")
        
        # Add the ImgBB URL to the response dictionary if it was generated
        if isinstance(api_response_data, dict):
            api_response_data['imgbb_url'] = imgbb_url

        return api_response_data
    except requests.exceptions.RequestException as e:
        print(f"API request failed: {e}")
        return None