From c43d5affb37e6c20e0f276bdb2f3110c3fc4bbc7 Mon Sep 17 00:00:00 2001
From: tonydero <tony@derocchis.com>
Date: Sun, 11 Aug 2024 19:21:10 -0600
Subject: [PATCH] initial attempt at very basic pywebio app

---
 flux-pywebio.py       | 74 +++++++++++++++++++++++++++++++++++++++++++
 flux-schnell-local.py |  2 +-
 2 files changed, 75 insertions(+), 1 deletion(-)
 create mode 100644 flux-pywebio.py

diff --git a/flux-pywebio.py b/flux-pywebio.py
new file mode 100644
index 0000000..b0f61b3
--- /dev/null
+++ b/flux-pywebio.py
@@ -0,0 +1,74 @@
+from pywebio import start_server
+from pywebio.input import *
+from pywebio.output import *
+import pywebio.pin as pin
+import time
+from diffusers import FluxPipeline
+import torch
+from pathlib import Path
+import re
+from datetime import datetime
+
+
+prev_prompt="A cat wearing an old racing helmet and goggles holding a sign that says zoom zoom"
+
+ckpt_id = "./FLUX.1-schnell"
+
+# denoising
+pipe = FluxPipeline.from_pretrained(
+    ckpt_id,
+    torch_dtype=torch.bfloat16,
+    use_safetensors=True,
+)
+pipe.vae.enable_tiling()
+pipe.vae.enable_slicing()
+pipe.enable_sequential_cpu_offload() # offloads modules to CPU on a submodule level (rather than model level)
+
+def slugify(text):
+    # remove non-word characters and foreign characters
+    text = re.sub(r"[^\w\s]", "", text)
+    text = re.sub(r"\s+", "-", text)
+    return text
+
+def flux_run(prompt, height, width, num_images_per_prompt, num_inference_steps, max_sequence_length, dirpath):
+    # TODO add check to make sure there isn't already one running from another source (e.g. from phone when on computer)
+    output = pipe(
+        prompt,
+        height=height,
+        width=width,
+        num_images_per_prompt=num_images_per_prompt,
+        num_inference_steps=num_inference_steps,
+        max_sequence_length=max_sequence_length,
+        guidance_scale=0.0,
+    )
+    # print('Max mem allocated (GB) while denoising:', torch.cuda.max_memory_allocated() / (1024 ** 3))
+
+    # import matplotlib.pyplot as plt
+    # plt.imshow(image)
+    # image.save("./whitehenge.png")
+    # plt.show()
+    for idx, image in enumerate(output.images):
+        timestamp = datetime.now().strftime("%Y%m%d%-H%M%S")
+        image_name = f'{slugify(prompt[:64])}-{idx}-{timestamp}.png'
+        image_path = dirpath / image_name
+        image.save(image_path)
+
+def btn_click(btn_val):
+    DIR_NAME="/home/tonydero/remdirs/immich/"
+    dirpath = Path(DIR_NAME)
+    if btn_val == 'Generate':
+        with put_loading():
+            put_text("Generating images...")
+            # time.sleep(3)  # Some time-consuming operations
+            flux_run(pin.pin.prompt, 720, 1280, 4, 4, 128, dirpath)
+        # put_text("The answer of the universe is 42")
+        toast("Images Generated!")
+
+def main():  # PyWebIO application function
+    put_text("Prompt:", inline=True)
+    pin.put_input("prompt", type="text", value=prev_prompt)
+    # TODO add the other parameters
+    put_buttons(['Generate'], onclick=btn_click)
+
+start_server(main, port=8080, debug=True, auto_open_webbrowser=False)
+
diff --git a/flux-schnell-local.py b/flux-schnell-local.py
index d40b2ef..308b7af 100644
--- a/flux-schnell-local.py
+++ b/flux-schnell-local.py
@@ -10,7 +10,7 @@ def slugify(text):
     text = re.sub(r"\s+", "-", text)
     return text
 
-prompt = "the town center of a small futuristic settlement with a fountain and gardens.the buildings are shaped like dodecahedrons, similar to small geodesic domes with a main road leading out to grassy plains sparsely dotted with very broad tall trees"
+prompt = "dirt highway road running through verdant grassy plains sparsely dotted with very broad tall trees with a mountain range and forest far in the distance"
 height, width = 720, 1280
 
 ckpt_id = "./FLUX.1-schnell"