utils/image.py

import sys
import os
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
from configs.config import Config

import base64, json
from openai import OpenAI
from datetime import datetime

client = OpenAI(
    base_url="https://api.studio.nebius.com/v1/",
    api_key=Config.nebius_api,
)

def generate_image_data(prompt: str, width: int = 1024, height: int = 1024, steps: int = 4, seed: int = -1, negative_prompt: str = "") -> dict:
    """
    Generates an image using the Nebius Studio model.

    Args:
        prompt (str): The prompt for image generation.
        width (int, optional): Width of the image. Default is 1024.
        height (int, optional): Height of the image. Default is 1024.
        steps (int, optional): Number of inference steps. Default is 4.
        seed (int, optional): Random seed. Default is -1 (random).
        negative_prompt (str, optional): Negative prompt to avoid unwanted features.

    Returns:
        dict: JSON response from the API, including base64 image.
    """
    response = client.images.generate(
        model="black-forest-labs/flux-schnell",
        response_format="b64_json",
        extra_body={
            "response_extension": "png",
            "width": width,
            "height": height,
            "num_inference_steps": steps,
            "negative_prompt": negative_prompt,
            "seed": seed
        },
        prompt=prompt
    )
    return response.to_dict()

def save_image_from_b64(image_data: str, output_folder: str = "image") -> str:
    """
    Decodes base64 image data and saves it as a PNG file.

    Args:
        image_data (str): Base64 encoded image string.
        output_folder (str): Folder where the image will be saved.

    Returns:
        str: Path to the saved image file.
    """
    # Create the output directory if it doesn't exist
    os.makedirs(output_folder, exist_ok=True)

    # Generate unique filename
    filename = datetime.now().strftime("%Y%m%d_%H%M%S") + ".png"
    file_path = os.path.join(output_folder, filename)

    # Decode and save image
    with open(file_path, "wb") as f:
        f.write(base64.b64decode(image_data))

    return file_path

def generate_images(prompt: str):
    """
    Generate images based on the script using Nebius Studio.
    This is a placeholder function that simulates image generation.
    """
    # For simplicity, we just return a mock image path
    image_data = generate_image_data(prompt)
    
    if "Error" in image_data:
        return {"error": image_data}
    
    # Save the generated image from base64 data
    image_path = save_image_from_b64(image_data['data'][0]['b64_json'])
    
    return {"image_path": image_path}