src/utils/generate_image.py

import os
import base64
from typing import Dict, Any
from openai import OpenAI

async def generate_image(
    prompt: str,
    output_path: str = "generated_image.png",
    width: int = 1024,
    height: int = 1024,
    num_inference_steps: int = 28,
    negative_prompt: str = "",
    seed: int = -1
) -> Dict[str, Any]:
    """
    Generate an image using Nebius API.
    
    Args:
        prompt: Text prompt for image generation
        output_path: Path where to save the generated image
        width: Image width
        height: Image height
        num_inference_steps: Number of inference steps
        negative_prompt: Negative prompt for generation
        seed: Random seed (-1 for random)
    
    Returns:
        Dictionary with result information
    """

    try:
        client = OpenAI(
            base_url="https://api.studio.nebius.com/v1/",
            api_key=os.environ.get("NEBIUS_API_KEY")
        )
        
        response = client.images.generate(
            model="black-forest-labs/flux-dev",
            response_format="b64_json",
            extra_body={
                "response_extension": "png",
                "width": width,
                "height": height,
                "num_inference_steps": num_inference_steps,
                "negative_prompt": negative_prompt,
                "seed": seed
            },
            prompt=prompt
        )
        
        image_data = base64.b64decode(response.data[0].b64_json)
        
        with open(output_path, 'wb') as output_file:
            output_file.write(image_data)
        
        output_size = os.path.getsize(output_path)
        
        return {
            "success": True,
            "message": "Image generated successfully",
            "prompt": prompt,
            "output_path": output_path,
            "output_size_bytes": output_size,
            "generation_params": {
                "width": width,
                "height": height,
                "num_inference_steps": num_inference_steps,
                "negative_prompt": negative_prompt,
                "seed": seed
            }
        }
        
    except Exception as e:
        if "NEBIUS_API_KEY" in str(e) or not os.environ.get("NEBIUS_API_KEY"):
            return {
                "success": False,
                "error": "NEBIUS_API_KEY environment variable not set",
                "output_path": None,
                'user': os.environ.get("USER")
            }
        return {
            "success": False,
            "error": f"Failed to generate image: {str(e)}",
            "output_path": None
        }