Update api.py

api.py

@@ -1,10 +1,14 @@
 from fastapi import FastAPI, UploadFile, File
 from fastapi.middleware.cors import CORSMiddleware
 from diffusers import AutoPipelineForText2Image
-import torch
 from PIL import Image
+import torch
 import io
+import base64
 from deep_translator import GoogleTranslator
+from rembg import remove
+import numpy as np
+import cv2
 
 app = FastAPI()
 
@@ -16,19 +20,46 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+print("Loading SDXL Turbo...")
+
 pipe = AutoPipelineForText2Image.from_pretrained(
     "stabilityai/sdxl-turbo",
-    torch_dtype=torch.float16,
-    variant="fp16"
+    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
 )
 
-pipe.to("cuda" if torch.cuda.is_available() else "cpu")
+pipe = pipe.to(device)
+
+
+def img_to_base64(img):
+    buf = io.BytesIO()
+    img.save(buf, format="PNG")
+    return base64.b64encode(buf.getvalue()).decode()
+
 
+def base64_to_img(data):
+    return Image.open(io.BytesIO(base64.b64decode(data)))
+
+
+def translate_prompt(prompt):
+    try:
+        return GoogleTranslator(source="auto", target="en").translate(prompt)
+    except:
+        return prompt
+
+
+# -----------------------
+# GENERATE IMAGE
+# -----------------------
 
 @app.post("/generate")
-async def generate(prompt: str):
 
-    prompt_en = GoogleTranslator(source='auto', target='en').translate(prompt)
+async def generate(data: dict):
+
+    prompt = data.get("prompt", "")
+
+    prompt_en = translate_prompt(prompt)
 
     image = pipe(
         prompt_en,
@@ -36,7 +67,97 @@ async def generate(prompt: str):
         guidance_scale=0.0
     ).images[0]
 
-    buf = io.BytesIO()
-    image.save(buf, format="PNG")
+    return {"image": img_to_base64(image)}
+
+
+# -----------------------
+# PRODUCT IMAGE
+# -----------------------
+
+@app.post("/product")
+
+async def product(data: dict):
+
+    prompt = data.get("prompt", "")
+
+    prompt_en = translate_prompt(prompt)
+
+    prompt_en += ", product photography, studio lighting, white background"
+
+    image = pipe(
+        prompt_en,
+        num_inference_steps=2,
+        guidance_scale=0.0
+    ).images[0]
+
+    return {"image": img_to_base64(image)}
+
+
+# -----------------------
+# UPSCALE
+# -----------------------
+
+@app.post("/upscale")
+
+async def upscale(file: UploadFile = File(...)):
+
+    image = Image.open(file.file).convert("RGB")
+
+    w, h = image.size
+
+    image = image.resize((w * 2, h * 2), Image.LANCZOS)
+
+    return {"image": img_to_base64(image)}
+
+
+# -----------------------
+# RESTORE
+# -----------------------
+
+@app.post("/restore")
+
+async def restore(file: UploadFile = File(...)):
+
+    image = Image.open(file.file).convert("RGB")
+
+    img = np.array(image)
+
+    img = cv2.fastNlMeansDenoisingColored(img, None, 10, 10, 7, 21)
+
+    result = Image.fromarray(img)
+
+    return {"image": img_to_base64(result)}
+
+
+# -----------------------
+# COLORIZE
+# -----------------------
+
+@app.post("/colorize")
+
+async def colorize(file: UploadFile = File(...)):
+
+    image = Image.open(file.file).convert("L")
+
+    img = np.array(image)
+
+    color = cv2.applyColorMap(img, cv2.COLORMAP_JET)
+
+    result = Image.fromarray(color)
+
+    return {"image": img_to_base64(result)}
+
+
+# -----------------------
+# REMOVE BACKGROUND
+# -----------------------
+
+@app.post("/removebg")
+
+async def removebg(file: UploadFile = File(...)):
+
+    image = Image.open(file.file)
+
+    result = remove(image)
 
-    return {"image": buf.getvalue()}
+    return {"image": img_to_base64(result)}