application/utils/image_generation.py

# application/utils/image_generation.py - NEW FILE
from huggingface_hub import InferenceClient
import os
import base64
from io import BytesIO

def generate_image(prompt: str) -> str:
    """
    Generates an image based on the given prompt using the Hugging Face Inference API.

    Args:
        prompt: The text prompt to generate the image from.

    Returns:
        The base64 encoded image data, or None if an error occurred.
    """
    try:
        api_key = os.environ.get('auth')
        if not api_key:
            raise ValueError("Hugging Face API key ('auth') not found in environment variables.")

        client = InferenceClient(api_key=api_key)
        image = client.text_to_image(
            prompt,
            model="black-forest-labs/FLUX.1-schnell"
        )

        # Convert PIL Image to base64 string
        buffered = BytesIO()
        image.save(buffered, format="PNG")  # Or "JPEG" if appropriate
        img_str = base64.b64encode(buffered.getvalue()).decode('utf-8')
        return img_str

    except Exception as e:
        print(f"Error in image generation: {e}")
        return None