use new OctoAI architecture based on OctoAI server

app.py

@@ -1,5 +1,4 @@
 import gradio as gr
-from PIL import Image
 import numpy as np
 
 from fair import text_to_image
@@ -7,12 +6,12 @@ from fair import text_to_image
 model_id = "runwayml/stable-diffusion-v1-5"
 device = "cuda"
 
+
 def infer(prompt):
-    image_file = text_to_image(prompt)
-    image = np.array(Image.open(image_file).convert('RGB'))
-    images = [image]
-    return images
-    
+    image = text_to_image(prompt)
+    image = np.array(image.convert('RGB'))
+    return [image]
+
     
 css = """
         .gradio-container {

fair.py

@@ -1,16 +1,27 @@
+import base64
 import json
+import logging
 import os
 import time
+from io import BytesIO
 from typing import List, BinaryIO
-import logging
+
+from PIL import Image
+from octoai.client import Client as OctoAiClient
+
 logger = logging.getLogger()
 
 import requests
 import tempfile
 
 SERVER_ADDRESS = "https://faircompute.com:8000/api/v1"
+ENDPOINT_ADDRESS = "http://dikobraz.mooo.com:5000"
+TARGET_NODE = "119eccba-2388-43c1-bdb9-02133049604c"
 # SERVER_ADDRESS = "http://localhost:8000/api/v1"
-DOCKER_IMAGE = "faircompute/stable-diffusion:pytorch-1.13.1-cu116"
+# ENDPOINT_ADDRESS = "http://localhost:5000"
+# TARGET_NODE = None
+
+DOCKER_IMAGE = "faircompute/diffusion-octo:latest"
 
 
 class FairApiClient:
@@ -49,7 +60,7 @@ class FairApiClient:
 
         return response
 
-    def put_job(self, image: str, command: List[str], input_files, output_files):
+    def put_job(self, image: str, command: List[str], ports: List[tuple[int, int]], input_files, output_files):
         url = f"{self.server_address}/jobs"
         data = {
             'type': 'V016',
@@ -57,12 +68,12 @@ class FairApiClient:
                 'type': 'V016',
                 'image': image,
                 'runtime': 'nvidia',
-                'ports': [],
+                'ports': [[{"port": host_port, "ip": 'null'}, {"port": container_port, "protocol": "Tcp"}] for (host_port, container_port) in ports],
                 'command': command,
             },
             'input_files': input_files,
             'output_files': output_files,
-            'target_node': None,
+            'target_node': TARGET_NODE,
         }
         response = self.put(url=url, data=data)
 
@@ -113,7 +124,39 @@ class FairApiClient:
         print(f"Failed to receive {path}")
 
 
-def text_to_image(text):
+class EndpointClient:
+    def infer(self, prompt):
+        client = OctoAiClient()
+
+        inputs = {"prompt": {"text": prompt}}
+        response = client.infer(endpoint_url=f"{ENDPOINT_ADDRESS}/infer", inputs=inputs)
+
+        image_b64 = response["output"]["image_b64"]
+        image_data = base64.b64decode(image_b64)
+        image_data = BytesIO(image_data)
+        image = Image.open(image_data)
+
+        return image
+
+
+class ServerNotReadyException(Exception):
+    pass
+
+
+def wait_for_server(retries, timeout):
+    for i in range(retries):
+        try:
+            r = requests.get(ENDPOINT_ADDRESS)
+            r.raise_for_status()
+            return
+        except (requests.exceptions.ConnectionError, requests.exceptions.HTTPError, requests.exceptions.Timeout):
+            logger.info("Server is not ready yet")
+            time.sleep(timeout)
+    else:
+        raise ServerNotReadyException("Failed to start the server")
+
+
+def start_server():
     # default credentials will work only for local server built in debug mode
     email = os.getenv('FAIRCOMPUTE_EMAIL', "debug-email")
     password = os.environ.get('FAIRCOMPUTE_PASSWORD', "debug-pwd")
@@ -123,52 +166,34 @@ def text_to_image(text):
     job_id = client.put_job(
         image=DOCKER_IMAGE,
         command=[],
+        ports=[(5000, 8080)],
         input_files=[],
-        output_files=["/workspace/result.png"])
+        output_files=[])
 
     logger.info(job_id)
 
     status = client.get_job_status(job_id=job_id)
     logger.info(status)
 
-    while status != "Processing" and status != "Completed":
+    while status != "Processing":
         status = client.get_job_status(job_id=job_id)
         logger.info(status)
         time.sleep(0.5)
 
-    res = client.put_job_stream_data(
-        job_id=job_id,
-        name="stdin",
-        data=text + "\n")
-    logger.info(res)
+    # wait until the server is ready
+    wait_for_server(retries=10, timeout=1.0)
 
-    res = client.put_job_stream_eof(
-        job_id=job_id,
-        name="stdin")
-    logger.info(res)
 
-    status = client.get_job_status(job_id=job_id)
-    logger.info(status)
+def text_to_image(text):
+    try:
+        wait_for_server(retries=1, timeout=0.0)
+    except ServerNotReadyException:
+        start_server()
 
-    while status == "Processing":
-        status = client.get_job_status(
-            job_id=job_id)
-        logger.info(status)
-        time.sleep(0.5)
-    if status == "Completed":
-        logger.info("Done!")
-    else:
-        logger.info("Job Failed")
-    file = client.wait_for_file(
-        job_id=job_id,
-        path="%2Fworkspace%2Fresult.png")
-    return file
+    client = EndpointClient()
+    return client.infer(text)
 
 
-if __name__=="__main__":
-    from PIL import Image
-    file = text_to_image(text="Robot dinozaur\n")
-    image = Image.open(file)
+if __name__ == "__main__":
+    image = text_to_image(text="Robot dinozaur\n")
     image.save("result.png")
-
-

requirements.txt

@@ -1 +1,2 @@
-gradio < 4
+gradio < 4
+octoai-sdk