initial attempt at very basic pywebio app

This commit is contained in:
tonydero 2024-08-11 19:21:10 -06:00
parent 2fca7f89f0
commit c43d5affb3
2 changed files with 75 additions and 1 deletions

74
flux-pywebio.py Normal file
View File

@ -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)

View File

@ -10,7 +10,7 @@ def slugify(text):
text = re.sub(r"\s+", "-", text) text = re.sub(r"\s+", "-", text)
return 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 height, width = 720, 1280
ckpt_id = "./FLUX.1-schnell" ckpt_id = "./FLUX.1-schnell"