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, artist=request.artist, website=request.website, resolution=request.resolution, additional_details=request.additional_details, color=request.color, 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, artist=request.artist, website=request.website, resolution=request.resolution, additional_details=request.additional_details, color=request.color, 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, artist=request.artist, website=request.website, resolution=request.resolution, additional_details=request.additional_details, color=request.color, lightning=request.lightning, ) image = decode_image(request.image) image.resize((256, 256)) if prompt: prompt = f"best quality, high quality, {prompt}" crud.create_prompt(db=db, user_id=current_user.user_id, prompt=prompt) else: request.scale = 1.0 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, negative_prompt=request.negative_prompt, scale=request.scale, guidance_scale=request.guidance_scale, ) images = [encode_image(image) for image in images] return images