Create utils/image.py

utils/image.py

@@ -0,0 +1,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}