Update app.py

app.py

@@ -29,10 +29,9 @@ def save_to_s3(image_data, payload, file_name):
     payload_data = io.BytesIO(payload.encode('utf-8'))
     s3.upload_fileobj(payload_data, bucket_name, payload_key, ExtraArgs={'ContentType': 'application/json'})
 
-
+import concurrent
 # Function to handle the NovelAI API request
-def generate_novelai_image(input_text, quality_tags, seed, negative_prompt, scale, ratio, sampler):
-    jwt_token = os.environ.get('NAI_API_KEY')
+def generate_novelai_image(input_text, quality_tags, seed, negative_prompt, scale, ratio, sampler, progress=gr.Progress()):
     if ratio == "Landscape (1216x832)":
         width = 1216
         height = 832
@@ -92,6 +91,20 @@ def generate_novelai_image(input_text, quality_tags, seed, negative_prompt, scal
         }
     }
 
+    futures = []
+    with concurrent.futures.ThreadPoolExecutor() as executor:
+        future = executor.submit(request_novelai, url, payload, headers)  
+        futures.append(future)
+
+        for _ in progress.tqdm([None] * 15, desc="Generating"):
+            done, not_done = concurrent.futures.wait(futures, timeout=np.random.uniform(0.8, 1.4))                 
+            if len(not_done) == 0:
+                break 
+        return futures[0].result()
+
+
+def request_novelai(url, payload, headers):
+
     # Send the POST request
     response = requests.post(url, json=payload, headers=headers)
 
@@ -108,7 +121,7 @@ def generate_novelai_image(input_text, quality_tags, seed, negative_prompt, scal
                     buffered = io.BytesIO()
                     image.save(buffered, format="WEBP", quality=98)
                     file_name = str(int(datetime.datetime.now().timestamp()))
-                    save_to_s3(buffered, json.dumps(payload, indent=4), file_name)
+                    # save_to_s3(buffered, json.dumps(payload, indent=4), file_name)
 
                     return np.array(image), json.dumps(payload, indent=4)
 
@@ -117,6 +130,7 @@ def generate_novelai_image(input_text, quality_tags, seed, negative_prompt, scal
     else:
         return "The response is not a zip file.", json.dumps(payload, indent=4)
 
+
 # Create Gradio interface
 iface = gr.Interface(
     fn=generate_novelai_image,