df 0.3 async requests test

app.py

@@ -14,9 +14,13 @@ import matplotlib.pyplot as plt
 import matplotlib.image as mpimg
 import time
 
+#----------------------------------------------------------------------------------------------------
+#async requests (otherwise each new !deepfloydif will add another ~20s until the first output)
+import aiohttp
 
 
 
+                
 #todos
 #alert
 
@@ -25,10 +29,31 @@ jojogan = gradio_client.Client("akhaliq/JoJoGAN")
 
 #token update
 DFIF_TOKEN = os.getenv('DFIF_TOKEN')
+hf_token = DFIF_TOKEN # could clean this up
 #deepfloydif
 #df = Client("DeepFloyd/IF") #not reliable at the moment
-df = Client("huggingface-projects/IF", hf_token=DFIF_TOKEN)
-
+#df = Client("huggingface-projects/IF", hf_token=DFIF_TOKEN) # testing replace with new async client
+df = AsyncGradioClient("huggingface-projects/IF", hf_token=DFIF_TOKEN)
+
+#----------------------------------------------------------------------------------------------------
+class AsyncGradioClient(gradio_client.Client):
+    async def predict_async(self, *args, api_name=None, fn_index=None):
+        if not self.is_live:
+            raise Exception("Cannot call predict() on a non-live interface. "
+                            "Please call launch() to launch the interface first.")
+        if api_name is None:
+            api_name = self.default_api_name
+        if fn_index is None:
+            fn_index = self.default_fn_index
+        url = self.api_url + api_name
+        headers = {"Authorization": "Bearer " + self.hf_token} if self.hf_token else {}
+        async with aiohttp.ClientSession() as session:
+            async with session.post(url, json={"args": args, "fn_index": fn_index}, headers=headers) as response:
+                response_json = await response.json()
+                if "error" in response_json:
+                    raise Exception(response_json["error"])
+                return response_json["output"]
+#----------------------------------------------------------------------------------------------------
 
 # Set up discord bot 
 class MyClient(discord.Client):
@@ -89,7 +114,8 @@ class MyClient(discord.Client):
                 # predict(img, model, api_name="/predict") -> output
                 # im = jojogan.predict(img, model) -> output
                 # im = jojogan.predict(attachment.url, style) -> output
-                
+
+                #im = await jojogan.predict_async(attachment.url, style)   # new async, could test
                 im = jojogan.predict(attachment.url, style)
                 await message.reply(f'Here is the {style} version of it', file=discord.File(im))
             else:
@@ -124,10 +150,13 @@ class MyClient(discord.Client):
             current_time = int(time.time())
             random.seed(current_time)
             seed = random.randint(1, 10000)
-            
-            stage_1_results, stage_1_param_path, stage_1_result_path = df.predict("llama", "blur", seed,1,7.0, 'smart100',50, api_name="/generate64")
 
             #stage_1_results, stage_1_param_path, stage_1_result_path = df.predict("gradio written on a wall", "blur", 1,1,7.0, 'smart100',50, api_name="/generate64")
+            stage_1_results, stage_1_param_path, stage_1_result_path = df.predict("llama", "blur", seed,1,7.0, 'smart100',50, api_name="/generate64")
+
+            stage_1_results, stage_1_param_path, stage_1_result_path = df.predict_async("llama", "blur", seed,1,7.0, 'smart100',50, api_name="/generate64")
+
+            
             
             # Assuming stage_1_results contains the path to the directory
             png_files = [f for f in os.listdir(stage_1_results) if f.endswith('.png')]