Update app.py

app.py

@@ -5,7 +5,7 @@ from typing import List, Optional, Union
 import os
 import numpy as np
 import torch
-
+import banana_dev as banana
 import PIL
 from diffusers import AutoencoderKL, DDIMScheduler, DiffusionPipeline, PNDMScheduler, UNet2DConditionModel
 from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker
@@ -13,41 +13,37 @@ from tqdm.auto import tqdm
 from transformers import CLIPFeatureExtractor, CLIPTextModel, CLIPTokenizer
 import gradio as gr
 import random
-import StableDiffusionInpaintingPipelineCustom
-
-device = "cuda"
-model_path = "CompVis/stable-diffusion-v1-4"
-
-StableDiffusionInpaintingPipeline = StableDiffusionInpaintingPipelineCustom.StableDiffusionInpaintingPipeline
-pipe = StableDiffusionInpaintingPipeline.from_pretrained(
-      model_path,
-      revision="fp16", 
-      torch_dtype=torch.float16,
-      use_auth_token=os.environ.get("access_token")).to(device)
-  
+import base64
+from io import BytesIO
+import os
 import gdown
+
 def download_gdrive_url():
   url = 'https://drive.google.com/u/0/uc?id=1PPO2MCttsmSqyB-vKh5C7SumwFKuhgyj&export=download'
   output = 'haarcascade_frontalface_default.xml'
   gdown.download(url, output, quiet=False)
   
-from torch import autocast
 def inpaint(p, init_image, mask_image=None, strength=0.75, guidance_scale=7.5, generator=None, num_samples=1, n_iter=1):
-    all_images = []
-    for _ in range(n_iter):
-        with autocast("cuda"):
-            images = pipe(
-                prompt=[p] * num_samples,
-                init_image=init_image,
-                mask_image=mask_image,
-                strength=strength,
-                guidance_scale=guidance_scale,
-                generator=generator,
-                num_inference_steps=75
-            )["sample"]
-        all_images.extend(images)
-    print(len(all_images))
-    return all_images[0]
+    buffered_init_img = BytesIO()
+    buffered_inverted_img = BytesIO()
+    init_image.save(buffered_init_img,format="JPEG")
+    mask_image.save(buffered_inverted_img,format="JPEG")
+    encoded_init_image = base64.b64encode(buffered_init_img.getvalue()).decode('utf-8')
+    encoded_inverted_image = base64.b64encode(buffered_inverted_img.getvalue()).decode('utf-8')
+    model_inputs = {
+      "prompt": "4K UHD professional profile picture of a person wearing a suit for work and posing for a picture, fine details, realistic shaded.",
+      "init_image": encoded_init_image,
+      "mask_image": encoded_inverted_image,
+      "strength": 0.65,
+      "guidance_scale": 10,
+      "num_inference_steps": 100
+    }
+    out = banana.run(os.environ.get("API_KEY"), os.environ.get("MODEL_KEY"), model_inputs)
+    image_byte_string = out["modelOutputs"][0]["output_image_base64"]
+    image_encoded = image_byte_string.encode('utf-8')
+    image_bytes = BytesIO(base64.b64decode(image_encoded))
+    return_image = Image.open(image_bytes)
+    return return_image
     
 def identify_face(user_image):
   img = cv2.imread(user_image.name) # read the resized image in cv2