From c8b41f61208699e3c465b8f0df156f03a617cc4f Mon Sep 17 00:00:00 2001
From: Tony <tony@rtd.one>
Date: Wed, 14 Aug 2024 20:23:28 -0600
Subject: [PATCH] fix prompt length check and move input labels

---
 flux-pywebio.py | 42 ++++++++++++++++++++++++------------------
 1 file changed, 24 insertions(+), 18 deletions(-)

diff --git a/flux-pywebio.py b/flux-pywebio.py
index 6a6813a..3ea1e45 100644
--- a/flux-pywebio.py
+++ b/flux-pywebio.py
@@ -46,7 +46,9 @@ def set_dir():
 
 def chk_prompt_len(prompt):
     if len(prompt) > 370:
-        return "Prompt too long"
+        return False
+    else:
+        return True
 
 def flux_run(prompt, height, width, num_images_per_prompt, num_inference_steps, max_sequence_length, dirpath):
     # FIX: add check to make sure there isn't already one running from another source (e.g. from phone when on computer)
@@ -71,15 +73,18 @@ def btn_click(btn_val):
     dirpath = set_dir()
     if dirpath:
         if btn_val == 'Generate':
-            with put_loading():
-                start_time = datetime.now()
-                put_text("Started generating images at " + start_time.strftime("%H:%M:%S"))
-                flux_run(pin.pin.prompt, pin.pin.height, pin.pin.width, pin.pin.numimgs, 4, 128, dirpath)
-            stop_time = datetime.now()
-            run_time = stop_time - start_time
-            # rt_min = run_time.minute
-            # rt_sec = run_time.second
-            put_success("Images finished generating in " + str(run_time), closable=True)
+            if chk_prompt_len(pin.pin.prompt):
+                with put_loading():
+                    start_time = datetime.now()
+                    put_text("Started generating images at " + start_time.strftime("%H:%M:%S"))
+                    flux_run(pin.pin.prompt, pin.pin.height, pin.pin.width, pin.pin.numimgs, 4, 128, dirpath)
+                stop_time = datetime.now()
+                run_time = stop_time - start_time
+                # rt_min = run_time.minute
+                # rt_sec = run_time.second
+                put_success("Images finished generating in " + str(run_time), closable=True)
+            else:
+                put_error("Prompt too long.", closable=True)
         elif btn_val == 'Portrait':
             pin.pin.width = 512
             pin.pin.height = 1024
@@ -94,15 +99,16 @@ def btn_click(btn_val):
 
 def main():  # PyWebIO application function
     # TODO: figure out how to put the labels in line with the inputs and make the number ones smaller
-    put_text("Prompt:", inline=True)
-    pin.put_input("prompt", type="text", value=default_prompt, validate=chk_prompt_len)
-    put_text("width:", inline=True)
-    pin.put_input("width", type="number", value=default_width)
-    put_text("height:", inline=True)
-    pin.put_input("height", type="number", value=default_height)
-    put_text("number of images:", inline=True)
-    pin.put_input("numimgs", type="number", value=default_numimgs)
+    # put_text("Prompt:", inline=True)
+    pin.put_input("prompt", label="Prompt:", type="text", value=default_prompt)
+    # put_text("width:", inline=True)
+    pin.put_input("width", label="Width:", type="number", value=default_width)
+    # put_text("height:", inline=True)
+    pin.put_input("height", label="Height:", type="number", value=default_height)
+    # put_text("number of images:", inline=True)
+    pin.put_input("numimgs", label="Number of images:", type="number", value=default_numimgs)
     # TODO: add the remaining parameters?
+    # TODO: add ability to loop for more generation in one run?
     put_buttons(['Generate', 'Portrait', 'Landscape', 'Square'], onclick=btn_click)
 
 start_server(main, port=4972, debug=True, auto_open_webbrowser=False)