Adldoke / main.py
Ashrafb's picture
Update main.py
f44bb3a verified
raw
history blame
No virus
1.42 kB
from fastapi import FastAPI, HTTPException
from fastapi import FastAPI
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse
import os
import requests
app = FastAPI()
API_URL = "https://ashrafb-dreamlikeart-diffusion-1-0.hf.space/"
HF_TOKEN = os.environ.get("HF_TOKEN", None)
def make_prediction(prompt, noise_level=0.0, fn_index=0):
headers={"Authorization": f"Bearer {HF_TOKEN}"}
data = {"prompt": prompt, "noise_level": noise_level}
response = requests.post(API_URL, headers=headers, json=data)
if response.status_code == 200:
return response.json()
else:
raise HTTPException(status_code=response.status_code, detail=response.text)
@app.get("/short-prompt/")
async def short_prompt(prompt: str):
try:
result = make_prediction(prompt)
return {"result": result}
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.get("/long-prompt/")
async def long_prompt(prompt: str, noise_level: float = 0.0):
try:
result = make_prediction(prompt, noise_level, fn_index=1)
return {"result": result}
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
app.mount("/", StaticFiles(directory="static", html=True), name="static")
@app.get("/")
def index() -> FileResponse:
return FileResponse(path="/app/static/index.html", media_type="text/html")