use just endpoint in text to image demo

text_to_image.py

@@ -1,158 +1,36 @@
 import base64
-import logging
-import os
-import hashlib
-
-import requests
-import time
 from io import BytesIO
 
+import requests
 from PIL import Image
-from fair import FairClient
-
-logger = logging.getLogger()
-
-SERVER_ADDRESS = "https://faircompute.com:8000"
-INFERENCE_NODE = "magnus"
-TUNNEL_NODE = "gcs-e2-micro"
-# SERVER_ADDRESS = "http://localhost:8000"
-# INFERENCE_NODE = "ef09913249aa40ecba7d0097f7622855"
-# TUNNEL_NODE = "c312e6c4788b00c73c287ab0445d3655"
-
-INFERENCE_DOCKER_IMAGE = "faircompute/diffusers-api-dreamshaper-8"
-TUNNEL_DOCKER_IMAGE = "rapiz1/rathole"
-
-endpoint_client = None
-fair_client = None
-
-
-class EndpointClient:
-    def __init__(self, server_address, timeout):
-        self.endpoint_address = f'http://{server_address}:5000'
-        response = requests.get(os.path.join(self.endpoint_address, 'healthcheck'), timeout=timeout).json()
-        if response['state'] != 'healthy':
-            raise Exception("Server is not healthy")
-
-    def infer(self, prompt):
-        inputs = {
-            "modelInputs": {
-                "prompt": prompt,
-                "num_inference_steps": 25,
-                "width": 512,
-                "height": 512,
-            },
-            "callInputs": {
-                "MODEL_ID": "lykon/dreamshaper-8",
-                "PIPELINE": "AutoPipelineForText2Image",
-                "SCHEDULER": "DEISMultistepScheduler",
-                "PRECISION": "fp16",
-                "REVISION": "fp16",
-            },
-        }
-
-        response = requests.post(self.endpoint_address, json=inputs).json()
-        image_data = BytesIO(base64.b64decode(response["image_base64"]))
-        image = Image.open(image_data)
-
-        return image
-
-
-class ServerNotReadyException(Exception):
-    pass
-
-
-def create_fair_client():
-    return FairClient(server_address=SERVER_ADDRESS,
-                      user_email=os.getenv('FAIRCOMPUTE_EMAIL', "debug-usr"),
-                      user_password=os.environ.get('FAIRCOMPUTE_PASSWORD', "debug-pwd"))
-
-
-def create_endpoint_client(fc, retries, timeout=1.0, delay=2.0):
-    nodes = fc.cluster().nodes.list()
-    server_address = next(info['host_address'] for info in nodes if info['name'] == TUNNEL_NODE)
-    for i in range(retries):
-        try:
-            return EndpointClient(server_address, timeout=timeout)
-        except (requests.exceptions.ConnectionError, requests.exceptions.HTTPError, requests.exceptions.Timeout) as e:
-            logging.exception(e)
-            time.sleep(delay)
-
-    raise ServerNotReadyException("Failed to start the server")
-
-
-def start_tunnel(fc: FairClient):
-    # generate fixed random authentication token based off some secret
-    token = hashlib.sha256(os.environ.get('FAIRCOMPUTE_PASSWORD', "debug-pwd").encode()).hexdigest()
-
-    # start tunnel node
-    server_config = f"""
-[server]
-bind_addr = "0.0.0.0:2333"  # port that rathole listens for clients
-
-[server.services.inference_server]
-token = "{token}"           # token that is used to authenticate the client for the service
-bind_addr = "0.0.0.0:5000"  # port that exposes service to the Internet
-"""
-    with open('server.toml', 'w') as file:
-        file.write(server_config)
-    fc.run(node_name=TUNNEL_NODE,
-           image=TUNNEL_DOCKER_IMAGE,
-           command=["--server", "/app/config.toml"],
-           volumes=[("./server.toml", "/app/config.toml")],
-           network="host",
-           detach=True)
-
-    nodes = fc.cluster().nodes.list()
-    server_address = next(info['host_address'] for info in nodes if info['name'] == TUNNEL_NODE)
-    client_config = f"""
-[client]
-remote_addr = "{server_address}:2333"       # address of the rathole server
-
-[client.services.inference_server]
-token = "{token}"                           # token that is used to authenticate the client for the service
-local_addr = "127.0.0.1:5001"               # address of the service that needs to be forwarded
-"""
-    with open('client.toml', 'w') as file:
-        file.write(client_config)
-    fc.run(node_name=INFERENCE_NODE,
-           image=TUNNEL_DOCKER_IMAGE,
-           command=["--client", "/app/config.toml"],
-           volumes=[("./client.toml", "/app/config.toml")],
-           network="host",
-           detach=True)
-
 
-def start_inference_server(fc: FairClient):
-    fc.run(node_name=INFERENCE_NODE,
-           image=INFERENCE_DOCKER_IMAGE,
-           runtime="nvidia",
-           ports=[(5001, 8000)],
-           detach=True)
+ENDPOINT_ADDRESS = "http://35.233.231.20:5000"
 
 
-def text_to_image(text):
-    global endpoint_client
-    global fair_client
-    if fair_client is None:
-        fair_client = create_fair_client()
+def text_to_image(prompt):
+    inputs = {
+        "modelInputs": {
+            "prompt": prompt,
+            "num_inference_steps": 25,
+            "width": 512,
+            "height": 512,
+        },
+        "callInputs": {
+            "MODEL_ID": "lykon/dreamshaper-8",
+            "PIPELINE": "AutoPipelineForText2Image",
+            "SCHEDULER": "DEISMultistepScheduler",
+            "PRECISION": "fp16",
+            "REVISION": "fp16",
+        },
+    }
 
-    try:
-        # client is configured, try to do inference right away
-        if endpoint_client is not None:
-            return endpoint_client.infer(text)
-        # client is not configured, try connecting to the inference server, maybe it is running
-        else:
-            endpoint_client = create_endpoint_client(fair_client, 1)
-    except (requests.exceptions.ConnectionError, requests.exceptions.HTTPError, requests.exceptions.Timeout, ServerNotReadyException):
-        # inference server is not ready, start inference server and open the tunnel
-        start_inference_server(fair_client)
-        start_tunnel(fair_client)
-        endpoint_client = create_endpoint_client(fair_client, retries=10)
+    response = requests.post(ENDPOINT_ADDRESS, json=inputs).json()
+    image_data = BytesIO(base64.b64decode(response["image_base64"]))
+    image = Image.open(image_data)
 
-    # run inference
-    return endpoint_client.infer(text)
+    return image
 
 
 if __name__ == "__main__":
-    image = text_to_image(text="Robot dinosaur\n")
+    image = text_to_image(prompt="Robot dinosaur")
     image.save("result.png")