endpoints.py

from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from routers import training
from huggingface_hub import login
from config import settings
import torch
# from diffusers import StableDiffusionPipeline
# from diffusers import StableDiffusionInpaintPipeline

# # pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", revision="fp16", torch_dtype=torch.float16)
# pipe = StableDiffusionInpaintPipeline.from_pretrained(
#     "stabilityai/stable-diffusion-2-inpainting",
#     torch_dtype=torch.float16,
# )

from diffusers import AutoPipelineForInpainting
from diffusers.utils import load_image, make_image_grid

pipee = AutoPipelineForInpainting.from_pretrained(
    "kandinsky-community/kandinsky-2-2-decoder-inpaint", 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"}