Update app.py

app.py

@@ -4,16 +4,36 @@ import requests
 from gradio_client import Client
 import base64
 
+from PIL import Image
+from io import BytesIO
+import base64
+import os
+
+from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, UniPCMultistepScheduler
+from diffusers.utils import load_image
+import torch
+
+import gradio as gr
+
+controlnet = ControlNetModel.from_pretrained("rgres/sd-controlnet-aerialdreams", torch_dtype=torch.float16)
+pipe = StableDiffusionControlNetPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-2-1-base", controlnet=controlnet, torch_dtype=torch.float16
+)
 
-base_gradio_url = os.getenv('URL_GRADIO', 'http://localhost:7860')
-client = None
+pipe = pipe.to("cuda")
+
+# CPU offloading for faster inference times
+pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
+pipe.enable_model_cpu_offload()
 
 app = Flask(__name__, static_url_path='/static')
 
+
 @app.route('/')
 def index():
     return app.send_static_file('index.html')
 
+
 def save_base64_image(base64Image):
     image_data = base64.b64decode(base64Image)
     path = "input_image.jpg"
@@ -21,11 +41,28 @@ def save_base64_image(base64Image):
         f.write(image_data)
     return path
 
+
 def encode_image_to_base64(filepath):
     with open(filepath, "rb") as image_file:
         encoded_image = base64.b64encode(image_file.read()).decode("utf-8")
     return encoded_image
 
+
+def generate_map(image, prompt, steps, seed):
+    #image = Image.open(BytesIO(base64.b64decode(image_base64)))
+    generator = torch.manual_seed(seed)
+
+    image = Image.fromarray(image)
+
+    image = pipe(
+        prompt=prompt,
+        num_inference_steps=steps,
+        image=image
+    ).images[0]
+
+    return image
+
+
 @app.route('/predict', methods=['POST'])
 def predict():
     global client
@@ -40,14 +77,9 @@ def predict():
     seed = data['data'][3]
 
     b64meta, b64_data = base64Image.split(',')
-
-    image_path = save_base64_image(b64_data)
-
-    result = client.predict(
-        image_path, prompt, steps, seed, fn_index=0
-    )
-
-    return b64meta + ',' + encode_image_to_base64(result)
+    image = Image.open(BytesIO(base64.b64decode(b64_data)))
+    
+    return (image, prompt, steps, seed)
 
 
 if __name__ == '__main__':