File size: 2,352 Bytes
d4174bb cc083b4 d4174bb cc083b4 d4174bb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 |
import os
import base64
from typing import Dict, Any
from openai import OpenAI
def generate_image(
prompt: str,
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)
return {
"success": True,
"message": "Image generated successfully",
"prompt": prompt,
"b64": image_data,
"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
} |