image_model.py

import requests
from pymongo import MongoClient
import os

uri = f"mongodb+srv://{os.getenv('mongo_secret')}@void-uep.guig8vk.mongodb.net/?retryWrites=true&w=majority"
client = MongoClient(uri)
db = client["ImagiGen"]
admin_collection = db["Admin"]

headers_flux = admin_collection.find_one({"model": "flux"})["key"] 
headers_flux = {"Authorization": f"Bearer {headers_flux}"}

def query_flux(API_URL, payload):
    response = requests.post(API_URL, headers=headers_flux, json=payload)
    return response.content


headers_stable_diffusion = admin_collection.find_one({"model": "stable"})["key"] 
headers_stable_diffusion = {"Authorization": f"Bearer {headers_stable_diffusion}"}
def query_satable_diffusion(API_URL, payload):
    response = requests.post(API_URL, headers=headers_stable_diffusion, 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_flux(
        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_satable_diffusion(
        API_URL,
        {
            "inputs": prompt,
        },
    )
    return image_bytes