image_models.py

import requests
import io
from PIL import Image
import os


headers = {"Authorization": f"Bearer {os.getenv('ImagiGen_HF_secret')}"}
def query(API_URL, payload):
    response = requests.post(API_URL, headers=headers, json=payload)
    return response.content

def flux(prompt):
    image = ""
    API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-dev"
    image_bytes = query(
        API_URL,
        {
            "inputs": prompt,
        },
    )
    return image_bytes

def stable_diffusion(prompt):
    API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-3-medium-diffusers"
    image_bytes = query(
        API_URL,
        {
            "inputs": prompt,
        },
    )
    return image_bytes