Update app/main.py

app/main.py

@@ -1,7 +1,75 @@
-from fastapi import FastAPI
+from fastapi import FastAPI, UploadFile, Form
+from fastapi.responses import StreamingResponse
+#import torch
+from PIL import Image
+#from diffusers import StableDiffusionDepth2ImgPipeline
+import numpy as np
+from io import BytesIO
 
 app = FastAPI()
 
-@app.get("/say_hello/")
-async def say_hello(name):
-    return {"message": f"Hello {name}!"}
+"""
+pipe = StableDiffusionDepth2ImgPipeline.from_pretrained(
+   "stabilityai/stable-diffusion-2-depth",
+   torch_dtype=torch.float16,
+).to("cuda")
+"""
+
+def pad_image(input_image):
+    pad_w, pad_h = np.max(((2, 2), np.ceil(
+        np.array(input_image.size) / 64).astype(int)), axis=0) * 64 - input_image.size
+    im_padded = Image.fromarray(
+        np.pad(np.array(input_image), ((0, pad_h), (0, pad_w), (0, 0)), mode='edge'))
+    w, h = im_padded.size
+    if w == h:
+        return im_padded
+    elif w > h:
+        new_image = Image.new(im_padded.mode, (w, w), (0, 0, 0))
+        new_image.paste(im_padded, (0, (w - h) // 2))
+        return new_image
+    else:
+        new_image = Image.new(im_padded.mode, (h, h), (0, 0, 0))
+        new_image.paste(im_padded, ((h - w) // 2, 0))
+        return new_image
+
+def predict(input_image, prompt, steps, scale, seed, strength, depth_image=None):
+    depth = None
+    if depth_image is not None:
+        depth_image = pad_image(depth_image)
+        depth_image = depth_image.resize((512, 512))
+        depth = np.array(depth_image.convert("L"))
+        depth = depth.astype(np.float32) / 255.0
+        depth = depth[None, None]
+        depth = torch.from_numpy(depth)
+    init_image = input_image.convert("RGB")
+    image = pad_image(init_image)  # resize to integer multiple of 32
+    image = image.resize((512, 512))
+    result = pipe(prompt=prompt, image=image, strength=strength)
+
+    return result['images']
+
+def grayscale(image,
+              prompt,
+              steps,
+              scale,
+              seed,
+              strength):
+    image = image.convert('L') #convert to grayscale
+
+    return image
+
+@app.post("/convert_ifc_img/")
+async def convert_ifc_img(file: UploadFile, 
+                          prompt: str = Form(default=""), 
+                          steps: int = Form(default=50), 
+                          scale: float = Form(default=9), 
+                          seed: int = Form(default=178106186), 
+                          strength: float = Form(default=0.9)
+                        ):
+    
+    image = Image.open(file.file)                        
+    image_result = grayscale(image, prompt, steps, scale, seed, strength)
+    buffer = BytesIO()
+    image_result.save(buffer, format="PNG")
+    buffer.seek(0)
+    return StreamingResponse(buffer, media_type="image/png")