Start Application

.gitignore

@@ -0,0 +1,160 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/

Dockerfile

@@ -0,0 +1,11 @@
+FROM python:3.11
+
+WORKDIR /code
+
+COPY ./requirements.txt /code/requirements.txt
+
+RUN pip install --no-cache-dir --upgrade -r /code/requirements.txt
+
+COPY . .
+
+CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "7860"]

app/api/generate.py

@@ -0,0 +1,159 @@
+from app.db import get_db
+from app.config import DEVICE
+from app.core import schemas, crud
+from app.security import get_current_user
+from app.core.schemas import TextImage, ImageImage, BackgroundRemoval, ImageVariations
+
+import base64
+from io import BytesIO
+from sqlalchemy.orm import Session
+from typing import Annotated, List
+from fastapi import APIRouter, Depends, HTTPException, Request
+
+import torch
+import numpy as np
+from PIL import Image
+import torch.nn.functional as F
+from torchvision.transforms.functional import normalize
+
+
+router = APIRouter()
+
+
+def decode_image(image):
+    return Image.open(BytesIO(base64.b64decode(image))).convert("RGB")
+
+
+def encode_image(image):
+    bytes = BytesIO()
+    image.save(bytes, format="PNG")
+    return base64.b64encode(bytes.getvalue())
+
+
+def create_prompt(subject, medium, style, artist, website, resolution, additional_details, color, lightning):
+    if not subject:
+        return None
+    if medium:
+        subject = f"{medium} of {subject}"
+    if style:
+        subject = f"{subject}, {style}"
+    if artist:
+        subject = f"{subject}, by {artist}"
+    if website:
+        subject = f"{subject}, {website}"
+    if resolution:
+        subject = f"{subject}, {resolution}"
+    if additional_details:
+        subject = f"{subject}, {additional_details}"
+    if color:
+        subject = f"{subject}, {color}"
+    if lightning:
+        subject = f"{subject}, {lightning}"
+    return subject
+
+
+@router.post("/text-image/", response_model=str)
+def text_image(model: Request, request: TextImage, db: Annotated[Session, Depends(get_db)], current_user: Annotated[schemas.User, Depends(get_current_user)]):
+    if not current_user.is_active:
+        raise HTTPException(status_code=403, detail="Forbidden")
+
+    generator = torch.manual_seed(request.seed)
+    prompt = create_prompt(request.prompt, medium=request.medium, style=request.style,
+                           additional_details=request.additional_details, lightning=request.lightning)
+
+    crud.create_prompt(db=db, user_id=current_user.user_id, prompt=prompt)
+
+    image = model.state.ti_pipe(prompt, num_inference_steps=request.num_inference_steps,
+                                guidance_scale=request.guidance_scale, generator=generator, negative_prompt=request.negative_prompt).images[0]
+
+    return encode_image(image)
+
+
+@router.post("/image-image/", response_model=str)
+def image_image(model: Request, request: ImageImage, db: Annotated[Session, Depends(get_db)], current_user: Annotated[schemas.User, Depends(get_current_user)]):
+    if not current_user.is_active:
+        raise HTTPException(status_code=403, detail="Forbidden")
+
+    generator = torch.manual_seed(request.seed)
+    prompt = create_prompt(request.prompt, medium=request.medium, style=request.style,
+                           additional_details=request.additional_details, lightning=request.lightning)
+    image = decode_image(request.image)
+
+    crud.create_prompt(db=db, user_id=current_user.user_id, prompt=prompt)
+
+    image = model.state.ii_pipe(prompt, image=image, num_inference_steps=request.num_inference_steps, guidance_scale=request.guidance_scale,
+                                image_guidance_scale=request.image_guidance_scale, generator=generator, negative_prompt=request.negative_prompt).images[0]
+
+    return encode_image(image)
+
+
+@router.post("/background-removal/", response_model=str)
+def background_removal(model: Request, request: BackgroundRemoval, current_user: Annotated[schemas.User, Depends(get_current_user)]):
+    if not current_user.is_active:
+        raise HTTPException(status_code=403, detail="Forbidden")
+
+    image = decode_image(request.image)
+
+    def preprocess_image(im: np.ndarray, model_input_size: list) -> torch.Tensor:
+        if len(im.shape) < 3:
+            im = im[:, :, np.newaxis]
+        # orig_im_size=im.shape[0:2]
+        im_tensor = torch.tensor(im, dtype=torch.float32).permute(2, 0, 1)
+        im_tensor = F.interpolate(torch.unsqueeze(
+            im_tensor, 0), size=model_input_size, mode='bilinear')
+        image = torch.divide(im_tensor, 255.0)
+        image = normalize(image, [0.5, 0.5, 0.5], [1.0, 1.0, 1.0])
+        return image
+
+    def postprocess_image(result: torch.Tensor, im_size: list) -> np.ndarray:
+        result = torch.squeeze(F.interpolate(
+            result, size=im_size, mode='bilinear'), 0)
+        ma = torch.max(result)
+        mi = torch.min(result)
+        result = (result-mi)/(ma-mi)
+        im_array = (result*255).permute(1, 2,
+                                        0).cpu().data.numpy().astype(np.uint8)
+        im_array = np.squeeze(im_array)
+        return im_array
+
+    # prepare input
+    model_input_size = [1024, 1024]
+    orig_im = np.array(image)
+    orig_im_size = orig_im.shape[0:2]
+    image = preprocess_image(orig_im, model_input_size).to(DEVICE)
+
+    # inference
+    result = model.state.br_model(image)
+
+    # post process
+    result_image = postprocess_image(result[0][0], orig_im_size)
+
+    # save result
+    pil_im = Image.fromarray(result_image)
+    no_bg_image = Image.new("RGBA", pil_im.size, (0, 0, 0, 0))
+    orig_image = Image.fromarray(orig_im)
+    no_bg_image.paste(orig_image, mask=pil_im)
+
+    return encode_image(no_bg_image)
+
+
+@router.post("/image-variations/", response_model=List[str])
+def image_variations(model: Request, request: ImageVariations, db: Annotated[Session, Depends(get_db)], current_user: Annotated[schemas.User, Depends(get_current_user)]):
+    if not current_user.is_active:
+        raise HTTPException(status_code=403, detail="Forbidden")
+
+    # prompt = create_prompt(request.prompt, medium=request.medium, style=request.style,
+    #                        additional_details=request.additional_details, lightning=request.lightning)
+    # image = decode_image(request.image)
+    # image.resize((512, 512))
+
+    # if prompt:
+    #     crud.create_prompt(db=db, user_id=current_user.user_id, prompt=prompt)
+
+    # images = model.state.iv_model.generate(pil_image=image, num_samples=request.num_samples, num_inference_steps=request.num_inference_steps,
+    #                                        seed=request.seed, prompt=prompt, scale=request.scale, negative_prompt=request.negative_prompt)
+
+    # images = [encode_image(image) for image in images]
+
+    # return images
+    return ["Image Variations is not supported yet."]

app/api/prompt.py

@@ -0,0 +1,34 @@
+from typing import Annotated, List
+from sqlalchemy.orm import Session
+from fastapi import APIRouter, Depends, HTTPException
+
+from app.db import get_db
+from app.core import schemas, crud
+from app.security import get_current_user
+
+
+router = APIRouter()
+
+
+@router.post("/get-all-prompts/", response_model=List[schemas.Prompt])
+def get_all_prompts(db: Annotated[Session, Depends(get_db)], current_user: Annotated[schemas.User, Depends(get_current_user)]):
+    if not current_user.is_superuser:
+        raise HTTPException(status_code=403, detail="Forbidden")
+
+    return crud.get_all_prompts(db=db)
+
+
+@ router.post("/get-prompt_by_user_id/{user_id}/", response_model=List[schemas.Prompt])
+def get_prompt_by_user_id(user_id: int, db: Annotated[Session, Depends(get_db)], current_user: Annotated[schemas.User, Depends(get_current_user)]):
+    if not current_user.is_superuser:
+        raise HTTPException(status_code=403, detail="Forbidden")
+
+    return crud.get_prompt_by_user_id(user_id=user_id, db=db)
+
+
+# @ router.post("/create-prompt/", response_model=schemas.Prompt)
+# def create_prompt(prompt: str, db: Annotated[Session, Depends(get_db)], current_user: Annotated[schemas.User, Depends(get_current_user)]):
+#     if not current_user.is_superuser:
+#         raise HTTPException(status_code=403, detail="Forbidden")
+
+#     return crud.create_prompt(prompt=prompt, db=db)

app/api/user.py

@@ -0,0 +1,55 @@
+from typing import Annotated, List
+from sqlalchemy.orm import Session
+from fastapi import APIRouter, Depends, HTTPException
+
+from app.db import get_db
+from app.core import schemas, crud
+from app.security import get_current_user
+
+
+router = APIRouter()
+
+
+@router.post("/create-user/", response_model=schemas.User)
+def create_user(user: schemas.UserCreate, db: Annotated[Session, Depends(get_db)], current_user: Annotated[schemas.User, Depends(get_current_user)]):
+    if not current_user.is_superuser:
+        raise HTTPException(status_code=403, detail="Forbidden")
+
+    user_exists = crud.get_user_by_username(username=user.username, db=db)
+    if user_exists:
+        raise HTTPException(
+            status_code=400, detail="Username already registered")
+
+    return crud.create_user(user=user, db=db)
+
+
+@router.post("/update-user/", response_model=schemas.User)
+def update_user(user: schemas.UserUpdate, db: Annotated[Session, Depends(get_db)], current_user: Annotated[schemas.User, Depends(get_current_user)]):
+    if not current_user.is_superuser:
+        raise HTTPException(status_code=403, detail="Forbidden")
+
+    user_exists = crud.get_user_by_user_id(user_id=user.user_id, db=db)
+    if not user_exists:
+        raise HTTPException(status_code=404, detail="User not found")
+
+    return crud.update_user(user=user, db=db)
+
+
+@router.post("/get-all-users/", response_model=List[schemas.User])
+def get_all_users(db: Annotated[Session, Depends(get_db)], current_user: Annotated[schemas.User, Depends(get_current_user)]):
+    if not current_user.is_superuser:
+        raise HTTPException(status_code=403, detail="Forbidden")
+
+    return crud.get_all_users(db=db)
+
+
+@router.post("/get-user_by_user_id/{user_id}/", response_model=schemas.User)
+def get_user_by_user_id(user_id: int, db: Annotated[Session, Depends(get_db)], current_user: Annotated[schemas.User, Depends(get_current_user)]):
+    if not current_user.is_superuser:
+        raise HTTPException(status_code=403, detail="Forbidden")
+
+    user = crud.get_user_by_user_id(user_id=user_id, db=db)
+    if user is None:
+        raise HTTPException(status_code=404, detail="User not found")
+
+    return user

app/config.py

@@ -0,0 +1,10 @@
+DATABASE_URL = 'sqlite:///./sql_app.db'
+
+SECRET_KEY = "20c7fc21a12dc8c19b06e603af10ebc11887a614f188f3d0878bb42bd784c315"
+ALGORITHM = "HS256"
+ACCESS_TOKEN_EXPIRE_MINUTES = 30
+
+SUPERUSER_USERNAME = "admin"
+SUPERUSER_PASSWORD = "admin"
+
+DEVICE = "cuda"

app/core/crud.py

@@ -0,0 +1,58 @@
+import app.sql.models as models
+import app.core.schemas as schemas
+
+from sqlalchemy.orm import Session
+from passlib.context import CryptContext
+
+pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
+
+
+def hash_password(password: str):
+    return pwd_context.hash(password)
+
+
+def create_prompt(prompt: str, user_id: int, db: Session):
+    db_prompt = models.Prompt(prompt=prompt, user_id=user_id)
+    db.add(db_prompt)
+    db.commit()
+    db.refresh(db_prompt)
+    return db_prompt
+
+
+def get_all_prompts(db: Session):
+    return db.query(models.Prompt).all()
+
+
+def get_prompt_by_user_id(user_id: int, db: Session):
+    return db.query(models.Prompt).filter(models.Prompt.user_id == user_id).all()
+
+
+def create_user(user: schemas.UserCreate, db: Session):
+    hashed_password = hash_password(user.password)
+    db_user = models.User(username=user.username, password=hashed_password)
+    db.add(db_user)
+    db.commit()
+    db.refresh(db_user)
+    return db_user
+
+
+def update_user(user: schemas.UserUpdate, db: Session):
+    db_user = db.query(models.User).filter(
+        models.User.user_id == user.user_id).first()
+    db_user.is_active = user.is_active
+    db_user.is_superuser = user.is_superuser
+    db.commit()
+    db.refresh(db_user)
+    return db_user
+
+
+def get_all_users(db: Session):
+    return db.query(models.User).all()
+
+
+def get_user_by_user_id(user_id: int, db: Session):
+    return db.query(models.User).filter(models.User.user_id == user_id).first()
+
+
+def get_user_by_username(username: str, db: Session):
+    return db.query(models.User).filter(models.User.username == username).first()

app/core/schemas.py

@@ -0,0 +1,77 @@
+from datetime import datetime
+from pydantic import BaseModel
+
+
+class PromptBase(BaseModel):
+    prompt: str
+
+
+class Prompt(BaseModel):
+    id: int
+    created_at: datetime
+
+    class Config:
+        from_attributes = True
+
+
+class UserBase(BaseModel):
+    pass
+
+
+class UserCreate(UserBase):
+    username: str
+    password: str
+
+
+class UserUpdate(UserBase):
+    user_id: int
+    is_active: bool = True
+    is_superuser: bool = False
+
+
+class User(UserBase):
+    user_id: int
+    username: str
+    is_active: bool
+    is_superuser: bool
+    created_at: datetime
+    updated_at: datetime
+    prompts: list[Prompt] = []
+
+    class Config:
+        from_attributes = True
+
+
+class Generate(BaseModel):
+    seed: int | None = None
+    negative_prompt : str = "ugly, tiling, poorly drawn hands, poorly drawn feet, poorly drawn face, out of frame, extra limbs, disfigured, deformed, body out of frame, bad anatomy, watermark, signature, cut off, low contrast, underexposed, overexposed, bad art, beginner, amateur, distorted face"
+    medium: str | None = None
+    style: str | None = None
+    lightning: str | None = None
+    additional_details: str | None = None
+
+
+class TextImage(Generate):
+    prompt: str
+    num_inference_steps: int = 4
+    guidance_scale: float = 2.0
+
+
+class ImageImage(Generate):
+    prompt: str
+    image: str
+    num_inference_steps: int = 10
+    guidance_scale: float = 7.5
+    image_guidance_scale: float = 1.5
+
+
+class BackgroundRemoval(BaseModel):
+    image: str
+
+
+class ImageVariations(Generate):
+    image: str
+    num_samples: int = 2
+    num_inference_steps: int = 30
+    prompt: str | None = None
+    scale: float = 0.5

app/db.py

@@ -0,0 +1,27 @@
+from app.sql.models import Base
+from app.sql.models import User
+from app.core.crud import hash_password
+from app.sql.database import SessionLocal, engine
+from app.config import SUPERUSER_USERNAME, SUPERUSER_PASSWORD
+
+
+def get_db():
+    db = SessionLocal()
+    try:
+        yield db
+    finally:
+        db.close()
+
+
+def init_db():
+    db = SessionLocal()
+    Base.metadata.create_all(bind=engine)
+
+    hashed_superuser_password = hash_password(SUPERUSER_PASSWORD)
+    superuser = User(username=SUPERUSER_USERNAME,
+                     password=hashed_superuser_password, is_superuser=True, is_active=True)
+
+    if db.query(User).filter(User.username == SUPERUSER_USERNAME).first() is None:
+        db.add(superuser)
+        db.commit()
+        db.refresh(superuser)

app/main.py

@@ -0,0 +1,101 @@
+from typing import Annotated
+from sqlalchemy.orm import Session
+from contextlib import asynccontextmanager
+from starlette.middleware.cors import CORSMiddleware
+from fastapi.security import OAuth2PasswordRequestForm
+from fastapi import APIRouter, FastAPI, HTTPException, Depends
+
+import torch
+from ip_adapter import IPAdapterXL
+from transformers import AutoModelForImageSegmentation
+from diffusers import AutoPipelineForText2Image, DPMSolverMultistepScheduler, StableDiffusionInstructPix2PixPipeline, EulerAncestralDiscreteScheduler, StableDiffusionXLPipeline
+
+from app.db import get_db, init_db
+from app.sql import models
+from app.api import user
+from app.api import prompt
+from app.sql.database import engine
+from app.api import generate
+from app.config import ACCESS_TOKEN_EXPIRE_MINUTES, DEVICE
+from app.security import authenticate_user, create_access_token, timedelta
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    base_model_path = "stabilityai/stable-diffusion-xl-base-1.0"
+    image_encoder_path = "sdxl_models/image_encoder"
+    ip_ckpt = "sdxl_models/ip-adapter_sdxl.bin"
+
+    ti_pipe = AutoPipelineForText2Image.from_pretrained(
+        'lykon/dreamshaper-xl-v2-turbo', torch_dtype=torch.float16, variant="fp16")
+    ti_pipe.to(DEVICE)
+    ti_pipe.scheduler = DPMSolverMultistepScheduler.from_config(
+        ti_pipe.scheduler.config)
+
+    ii_pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(
+        "timbrooks/instruct-pix2pix", torch_dtype=torch.float16, safety_checker=None)
+    ii_pipe.to(DEVICE)
+    ii_pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(
+        ii_pipe.scheduler.config)
+
+    br_model = AutoModelForImageSegmentation.from_pretrained(
+        "briaai/RMBG-1.4", trust_remote_code=True)
+    br_model.to(DEVICE)
+
+    # sdxl_pipe = StableDiffusionXLPipeline.from_pretrained(
+    #     base_model_path,
+    #     torch_dtype=torch.float16,
+    #     add_watermarker=False,
+    # )
+    # iv_model = IPAdapterXL(sdxl_pipe, image_encoder_path, ip_ckpt, DEVICE)
+
+    yield {'ti_pipe': ti_pipe, 'ii_pipe': ii_pipe, 'br_model': br_model}  # , 'iv_model': iv_model
+
+    del ti_pipe
+    del ii_pipe
+    del br_model
+    # del sdxl_pipe
+    # del iv_model
+
+
+app = FastAPI(lifespan=lifespan)
+# app = FastAPI()
+
+router = APIRouter()
+
+origins = ["*"]
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=origins,
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+init_db()
+
+
+@app.get("/")
+def read_root():
+    return {"Hello": "World"}
+
+
+@app.post("/token")
+async def login(form_data: Annotated[OAuth2PasswordRequestForm, Depends()], db: Annotated[Session, Depends(get_db)]):
+    user = authenticate_user(
+        db, form_data.username, form_data.password)
+    if not user:
+        raise HTTPException(
+            status_code=400, detail="Incorrect username or password")
+    access_token_expires = timedelta(
+        minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
+    access_token = create_access_token(
+        data={"sub": user.username}, expires_delta=access_token_expires)
+    return {"access_token": access_token, "token_type": "bearer"}
+
+
+router.include_router(user.router, prefix="/users")
+router.include_router(prompt.router, prefix="/prompts")
+router.include_router(generate.router, prefix="/generate")
+app.include_router(router)

app/security.py

@@ -0,0 +1,61 @@
+from typing import Annotated
+from datetime import datetime, timedelta, timezone
+from fastapi import Depends, HTTPException
+from fastapi.security import OAuth2PasswordBearer
+from passlib.context import CryptContext
+from sqlalchemy.orm import Session
+from jose import JWTError, jwt
+
+
+from app.config import ALGORITHM, SECRET_KEY
+from app.db import get_db
+from app.sql import models
+
+
+pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
+
+oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
+
+
+def verify_password(plain_password, hashed_password):
+    return pwd_context.verify(plain_password, hashed_password)
+
+
+def get_user(db: Session, username: str):
+    return db.query(models.User).filter(models.User.username == username).first()
+
+
+def authenticate_user(db: Session, username: str, password: str):
+    user = get_user(db, username)
+    if not user:
+        return False
+    if not verify_password(password, user.password):
+        return False
+    return user
+
+
+def create_access_token(data: dict, expires_delta: timedelta | None = None):
+    to_encode = data.copy()
+    if expires_delta:
+        expire = datetime.now(timezone.utc) + expires_delta
+    else:
+        expire = datetime.now(timezone.utc) + timedelta(minutes=15)
+    to_encode.update({"exp": expire})
+    encoded_jwt = jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)
+    return encoded_jwt
+
+
+async def get_current_user(token: Annotated[str, Depends(oauth2_scheme)], db: Annotated[Session, Depends(get_db)]):
+    credentials_exception = HTTPException(
+        status_code=401, detail="Could not validate credentials")
+    try:
+        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
+        username: str = payload.get("sub")
+        if username is None:
+            raise credentials_exception
+    except JWTError:
+        raise credentials_exception
+    user = get_user(db, username)
+    if user is None:
+        raise credentials_exception
+    return user

app/sql/database.py

@@ -0,0 +1,11 @@
+from app.config import DATABASE_URL
+from sqlalchemy import create_engine
+from sqlalchemy.orm import sessionmaker
+from sqlalchemy.ext.declarative import declarative_base
+
+
+engine = create_engine(DATABASE_URL, connect_args={"check_same_thread": False})
+
+SessionLocal = sessionmaker(bind=engine, autocommit=False, autoflush=False)
+
+Base = declarative_base()

app/sql/models.py

@@ -0,0 +1,29 @@
+from datetime import datetime
+from app.sql.database import Base
+from sqlalchemy.orm import relationship
+from sqlalchemy import Boolean, Column, DateTime, ForeignKey, Integer, String
+
+
+class User(Base):
+    __tablename__ = 'users'
+
+    user_id = Column(Integer, primary_key=True, index=True)
+    username = Column(String, index=True, unique=True)
+    password = Column(String)
+    is_superuser = Column(Boolean, default=False)
+    is_active = Column(Boolean, default=True)
+    created_at = Column(DateTime, default=datetime.now)
+    updated_at = Column(DateTime, default=datetime.now, onupdate=datetime.now)
+
+    prompts = relationship("Prompt", back_populates="user")
+
+
+class Prompt(Base):
+    __tablename__ = 'prompts'
+
+    id = Column(Integer, primary_key=True)
+    prompt = Column(String, index=True)
+    created_at = Column(DateTime, default=datetime.now)
+    user_id = Column(Integer, ForeignKey('users.user_id'))
+
+    user = relationship("User", back_populates="prompts")

requirements.txt

@@ -0,0 +1,7 @@
+fastapi
+uvicorn[standard]
+sqlalchemy
+psycopg2-binary
+python-multipart
+python-jose[cryptography]
+passlib[bcrypt]