Spaces:
Sleeping
Sleeping
File size: 1,376 Bytes
51acb88 78bdd91 51acb88 |
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 |
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
|