Spaces:
Runtime error
Runtime error
File size: 2,711 Bytes
e1012e2 |
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 80 81 82 83 |
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} |