74 lines
2.5 KiB
Python
74 lines
2.5 KiB
Python
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))
|
|
|
|
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():
|
|
start_time = datetime.now()
|
|
put_text("Started generating images at " + start_time.strftime("%H:%M:%S"))
|
|
flux_run(pin.pin.prompt, 720, 1280, 4, 4, 128, dirpath)
|
|
stop_time = datetime.now()
|
|
run_time = stop_time - start_time
|
|
# rt_min = run_time.minutes
|
|
# rt_sec = run_time.seconds
|
|
put_success("Images finished generating in " + str(run_time), closable=True)
|
|
|
|
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)
|
|
|