petrockapi / endpoints.py
woods-today
Working on it
6a05099
raw
history blame
870 Bytes
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from routers import inference, training
from huggingface_hub import login
from config import settings
import torch
from diffusers import StableDiffusionPipeline
pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", revision="fp16", torch_dtype=torch.float16)
login(settings.huggingface_key)
# app = FastAPI(openapi_url="/api/v1/sparrow-ml/openapi.json", docs_url="/api/v1/sparrow-ml/docs")
app = FastAPI()
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_methods=["*"],
allow_headers=["*"],
allow_credentials=True,
)
# app.include_router(inference.router, prefix="/api-inference/v1/sparrow-ml", tags=["Inference"])
app.include_router(training.router)
@app.get("/")
async def root():
return {"message": "Sparrow ML API"}