Add inference with fair compute api

app.py

@@ -1,36 +1,35 @@
 import gradio as gr
-from datasets import load_dataset
+#from datasets import load_dataset
 from PIL import Image  
 import re
 import os
 import requests
+import numpy as np
+
+from fair import text_to_image
 
 from share_btn import community_icon_html, loading_icon_html, share_js
 
 model_id = "runwayml/stable-diffusion-v1-5"
 device = "cuda"
 
-word_list_dataset = load_dataset("stabilityai/word-list", data_files="list.txt", use_auth_token=True)
-word_list = word_list_dataset["train"]['text']
+#word_list_dataset = load_dataset("stabilityai/word-list", data_files="list.txt", use_auth_token=True)
+#word_list = word_list_dataset["train"]['text']
 
-is_gpu_busy = False
+#is_gpu_busy = False
 def infer(prompt):
-    global is_gpu_busy
+    #global is_gpu_busy
     samples = 4
     steps = 50
     scale = 7.5
-    for filter in word_list:
-        if re.search(rf"\b{filter}\b", prompt):
-            raise gr.Error("Unsafe content found. Please try again with different prompts.")
-        
+    # for filter in word_list:
+    #     if re.search(rf"\b{filter}\b", prompt):
+    #         raise gr.Error("Unsafe content found. Please try again with different prompts.")
+    #
     images = []
-    url = os.getenv('JAX_BACKEND_URL')
-    payload = {'prompt': prompt}
-    images_request = requests.post(url, json = payload)
-    for image in images_request.json()["images"]:
-        image_b64 = (f"data:image/jpeg;base64,{image}")
-        images.append(image_b64)
-    
+    image = text_to_image(prompt)
+    image = np.array(Image.open(image).convert('RGB'))
+    images = [image, image, image, image]
     return images
     
     
@@ -239,55 +238,55 @@ with block:
                     rounded=(False, True, True, False),
                     full_width=False,
                 )
-
+#        gallery = gr.Image(type="filepath").style(grid=[2], height="auto")
         gallery = gr.Gallery(
             label="Generated images", show_label=False, elem_id="gallery"
         ).style(grid=[2], height="auto")
 
-        with gr.Group(elem_id="container-advanced-btns"):
-            advanced_button = gr.Button("Advanced options", elem_id="advanced-btn")
-            with gr.Group(elem_id="share-btn-container"):
-                community_icon = gr.HTML(community_icon_html)
-                loading_icon = gr.HTML(loading_icon_html)
-                share_button = gr.Button("Share to community", elem_id="share-btn")
+        # with gr.Group(elem_id="container-advanced-btns"):
+        #     advanced_button = gr.Button("Advanced options", elem_id="advanced-btn")
+        #     with gr.Group(elem_id="share-btn-container"):
+        #         community_icon = gr.HTML(community_icon_html)
+        #         loading_icon = gr.HTML(loading_icon_html)
+        #         share_button = gr.Button("Share to community", elem_id="share-btn")
+        #
+        # with gr.Row(elem_id="advanced-options"):
+        #     gr.Markdown("Advanced settings are temporarily unavailable")
+        #     samples = gr.Slider(label="Images", minimum=1, maximum=4, value=4, step=1)
+        #     steps = gr.Slider(label="Steps", minimum=1, maximum=50, value=45, step=1)
+        #     scale = gr.Slider(
+        #         label="Guidance Scale", minimum=0, maximum=50, value=7.5, step=0.1
+        #     )
+        #     seed = gr.Slider(
+        #         label="Seed",
+        #         minimum=0,
+        #         maximum=2147483647,
+        #         step=1,
+        #         randomize=True,
+        #     )
 
-        with gr.Row(elem_id="advanced-options"):
-            gr.Markdown("Advanced settings are temporarily unavailable")
-            samples = gr.Slider(label="Images", minimum=1, maximum=4, value=4, step=1)
-            steps = gr.Slider(label="Steps", minimum=1, maximum=50, value=45, step=1)
-            scale = gr.Slider(
-                label="Guidance Scale", minimum=0, maximum=50, value=7.5, step=0.1
-            )
-            seed = gr.Slider(
-                label="Seed",
-                minimum=0,
-                maximum=2147483647,
-                step=1,
-                randomize=True,
-            )
+        #ex = gr.Examples(examples=examples, fn=infer, inputs=text, outputs=[gallery], cache_examples=True, postprocess=False)
+        #ex.dataset.headers = [""]
 
-        ex = gr.Examples(examples=examples, fn=infer, inputs=text, outputs=[gallery], cache_examples=True, postprocess=False)
-        ex.dataset.headers = [""]
-
-        text.submit(infer, inputs=text, outputs=[gallery], postprocess=False)
-        btn.click(infer, inputs=text, outputs=[gallery], postprocess=False)
+        text.submit(infer, inputs=text, outputs=[gallery])
+        btn.click(infer, inputs=text, outputs=[gallery])
         
-        advanced_button.click(
-            None,
-            [],
-            text,
-            _js="""
-            () => {
-                const options = document.querySelector("body > gradio-app").querySelector("#advanced-options");
-                options.style.display = ["none", ""].includes(options.style.display) ? "flex" : "none";
-            }""",
-        )
-        share_button.click(
-            None,
-            [],
-            [],
-            _js=share_js,
-        )
+        # advanced_button.click(
+        #     None,
+        #     [],
+        #     text,
+        #     _js="""
+        #     () => {
+        #         const options = document.querySelector("body > gradio-app").querySelector("#advanced-options");
+        #         options.style.display = ["none", ""].includes(options.style.display) ? "flex" : "none";
+        #     }""",
+        # )
+        # share_button.click(
+        #     None,
+        #     [],
+        #     [],
+        #     _js=share_js,
+        # )
         gr.HTML(
             """
                 <div class="footer">

fair.py

@@ -0,0 +1,254 @@
+import json
+import os
+import time
+from typing import List
+import logging
+logger = logging.getLogger()
+
+import requests
+
+#SERVER_ADRESS="https://faircompute.com:8000/api/v1"
+SERVER_ADRESS="http://localhost:8000/api/v1"
+DOCKER_IMAGE="faircompute/stable-diffusion:pytorch-1.13.1-cu116"
+#DOCKER_IMAGE="sha256:e06453fe869556ea3e63572a935aed4261337b261fdf7bda370472b0587409a9"
+
+def authenticate(email: str, password: str):
+    url = f'{SERVER_ADRESS}/auth/login'
+    json_obj = {"email": email, "password": password}
+    resp = requests.post(url, json=json_obj)
+    token = resp.json()["token"]
+    return token
+
+def get(url, token, **kwargs):
+    headers = {
+        'Authorization': f'Bearer {token}'
+    }
+    response = requests.get(url, headers=headers, **kwargs)
+
+    if not response.ok:
+        raise Exception(f"Error! status: {response.status_code}")
+    return response
+
+
+def put(url, token, data):
+    headers = {
+        'Content-Type': 'application/json',
+        'Authorization': f'Bearer {token}'
+    }
+    if not isinstance(data, str):
+        data = json.dumps(data)
+    response = requests.put(url, headers=headers, data=data)
+
+    if not response.ok and response.status_code != 206:
+        raise Exception(f"Error! status: {response.status_code}")
+    return response
+
+
+def put_program(token, launcher: str, image: str, runtime: str, command: List[str]):
+    url = f"{SERVER_ADRESS}/programs"
+    data = {
+        launcher: {
+            "image": image,
+            "command": command,
+            "runtime": runtime
+        }
+    }
+    response = put(url=url, token=token, data=data)
+
+    return int(response.text)
+
+
+def put_job(token, program_id, input_files, output_files):
+    url = f"{SERVER_ADRESS}/jobs?program={program_id}"
+    data = {
+        'input_files': input_files,
+        'output_files': output_files
+    }
+
+    response = put(url=url, token=token, data=data)
+
+    return int(response.text)
+
+
+def get_job_status(token, job_id):
+    url = f"{SERVER_ADRESS}/jobs/{job_id}/status"
+    response = get(url=url, token=token)
+    return response.text
+
+
+def get_cluster_summary(token):
+    url = f"{SERVER_ADRESS}/nodes/summary"
+
+    response = get(token=token, url=url)
+
+    return response.json()
+
+
+def put_job_stream_data(token, job_id, name, data):
+    url = f"{SERVER_ADRESS}/jobs/{job_id}/data/streams/{name}"
+    response = put(url=url, token=token, data=data)
+
+    return response.text
+
+
+def put_job_stream_eof(token, job_id, name):
+    url = f"{SERVER_ADRESS}/jobs/{job_id}/data/streams/{name}/eof"
+
+    response = put(url=url, token=token, data=None)
+
+    return response.text
+
+
+def wait_for_file(token, job_id, path, local_path, attempts=10):
+    headers = {
+        'Authorization': f'Bearer {token}'
+    }
+    for i in range(attempts):
+        url = f"{SERVER_ADRESS}/jobs/{job_id}/data/files/{path}"
+        print(f"Waiting for file {path}...")
+        try:
+            with requests.get(url=url, headers=headers, stream=True) as r:
+                r.raise_for_status()
+                with open(local_path, 'wb') as f:
+                    for chunk in r.iter_content(chunk_size=8192):
+                        f.write(chunk)
+
+                print(f"File {local_path} ready")
+                return local_path
+        except Exception as e:
+            print(e)
+            time.sleep(0.5)
+
+    print(f"Failed to receive {local_path}")
+
+
+def text_to_image(text):
+    email = os.getenv('FAIRCOMPUTE_EMAIL')
+    password = os.environ.get('FAIRCOMPUTE_PASSWORD')
+    token = authenticate(email=email, password=password)
+
+    logger.info(token)
+
+    summary = get_cluster_summary(token=token)
+    logger.info("Summary:")
+    logger.info(summary)
+    program_id = put_program(token=token,
+                             launcher="Docker",
+                             image=DOCKER_IMAGE,
+                             runtime="nvidia",
+                             command=[])
+    logger.info(program_id)
+
+    job_id = put_job(token=token,
+                     program_id=program_id,
+                     input_files=[],
+                     output_files=["/workspace/result.png"])
+
+    logger.info(job_id)
+
+    status = get_job_status(token=token,
+                            job_id=job_id)
+    logger.info(status)
+
+    while status != "Processing" and status != "Completed":
+        status = get_job_status(token=token,
+                                job_id=job_id)
+        logger.info(status)
+        time.sleep(0.5)
+
+    res = put_job_stream_data(token=token,
+                              job_id=job_id,
+                              name="stdin",
+                              data=text + "\n")
+    logger.info(res)
+
+    res = put_job_stream_eof(token=token,
+                             job_id=job_id,
+                             name="stdin")
+    logger.info(res)
+
+    status = get_job_status(token=token,
+                            job_id=job_id)
+    logger.info(status)
+
+    while status == "Processing":
+        status = get_job_status(token=token,
+                                job_id=job_id)
+        logger.info(status)
+        time.sleep(0.5)
+    if status == "Completed":
+        logger.info("Done!")
+    else:
+        logger.info("Job Failed")
+    resp = wait_for_file(token=token,
+                         job_id=job_id,
+                         path="%2Fworkspace%2Fresult.png",
+                         local_path="result.png")
+    logger.info(resp)
+    return resp
+
+
+if __name__=="__main__":
+    email = os.getenv('FAIRCOMPUTE_EMAIL')
+    password = os.environ.get('FAIRCOMPUTE_PASSWORD')
+    token = authenticate(email=email, password=password)
+
+    print(token)
+
+    summary = get_cluster_summary(token=token)
+    print("Summary:")
+    print(summary)
+    program_id = put_program(token=token,
+                             launcher="Docker",
+                             image=DOCKER_IMAGE,
+                             runtime="nvidia",
+                             command=[])
+    print(program_id)
+
+    job_id = put_job(token=token,
+                     program_id=program_id,
+                     input_files=[],
+                     output_files=["/workspace/result.png"])
+
+    print(job_id)
+
+    status = get_job_status(token=token,
+                            job_id=job_id)
+    print(status)
+
+    while status != "Processing" and status != "Completed":
+         status = get_job_status(token=token,
+                                 job_id=job_id)
+         print(status)
+         time.sleep(0.5)
+
+    res = put_job_stream_data(token=token,
+                        job_id=job_id,
+                        name="stdin",
+                        data="Robot dinozaur\n")
+    print(res)
+
+    res = put_job_stream_eof(token=token,
+                       job_id=job_id,
+                       name="stdin")
+    print(res)
+
+    status = get_job_status(token=token,
+                            job_id=job_id)
+    print(status)
+
+    while status == "Processing":
+         status = get_job_status(token=token,
+                                 job_id=job_id)
+         print(status)
+         time.sleep(0.5)
+    if status == "Completed":
+        print("Done!")
+    else:
+        print("Job Failed")
+    resp = wait_for_file(token=token,
+                         job_id=job_id,
+                         path="%2Fworkspace%2Fresult.png",
+                         local_path="result.png")
+    print(resp)
+

requirements.txt

@@ -1 +1 @@
-python-dotenv
+gradio < 4