from collections import namedtuple import gradio as gr import lupa import numpy as np import os from pathlib import Path from PIL import Image import re import torch from torchvision import transforms import traceback from modules.call_queue import wrap_gradio_gpu_call from modules import scripts, script_callbacks, devices, ui, shared, processing, sd_samplers, sd_samplers_common, paths from modules import prompt_parser, ui, face_restoration, deepbooru import modules.images as images from modules.shared import opts, cmd_opts, state from modules.processing import StableDiffusionProcessingTxt2Img, Processed, process_images, fix_seed, decode_first_stage, apply_overlay, apply_color_correction, create_infotext, create_random_tensors try: from modules.rng import ImageRNG except ImportError: ImageRNG = None sd_webui_lua_dir = scripts.basedir() def filter_attribute_access(obj, attr_name, is_setting): if isinstance(attr_name, (str)): if not attr_name.startswith('_'): return attr_name raise AttributeError('access denied') L = lupa.LuaRuntime(register_eval=False, attribute_filter=filter_attribute_access) G = L.globals() LUA_output = '' LUA_gallery = [] def lua_run(id_task, lua_input, lua_code): global LUA_output, LUA_gallery print(f"LUA: START: {id_task}") # Crop --START-- and --END-- lua_code = re.sub('.*\n--START--', '', lua_code, 1, flags=re.S) lua_code = re.sub('\n--END--.*', '', lua_code, 1, flags=re.S) # Prepend input lua_code = f"input = [[{lua_input}]]\n{lua_code}" try: result = L.execute(lua_code) if result: LUA_output += str(result)+'\n' except Exception as err: traceback.print_exc() result = f"ERROR: {err}" print(f"LUA {result}") raise gr.Error(result) print(f"LUA: END: {id_task}") return LUA_output, LUA_gallery, '' def lua_reset(): global L, G, LUA_output, LUA_gallery L = lupa.LuaRuntime(register_eval=False, attribute_filter=filter_attribute_access) G = L.globals() LUA_output = '' LUA_gallery = [] # Setup python functions (messy list. Will most likely change) G.sd = { 'empty_latent': sd_lua_empty_latent, 'pipeline': sd_lua_pipeline, 'process': sd_lua_process, 'getp': sd_lua_getp, 'cond': sd_lua_cond, 'negcond': sd_lua_negcond, 'sample': sd_lua_sample, 'vae': sd_lua_vae, 'toimage': sd_lua_toimage, 'makegif': sd_lua_makegif, 'getsamplers': sd_lua_getsamplers, 'restorefaces': sd_lua_restorefaces, 'textencode': sd_lua_textencode, 'clip2negcond': sd_lua_clip2negcond, 'negcond2cond': sd_lua_negcond2cond, 'interrogate': { 'clip': sd_lua_interrogate_clip, 'blip': sd_lua_interrogate_clip, 'deepbooru': sd_lua_interrogate_deepbooru, } } G.ui = { 'clear': ui_lua_output_clear, 'console': ui_lua_console, 'out': ui_lua_output, 'gallery': { 'add': ui_lua_gallery_add, 'addc': ui_lua_gallery_addc, 'clear': ui_lua_gallery_clear, 'del': ui_lua_gallery_del, 'getgif': ui_lua_gallery_getgif, 'saveall': ui_lua_gallery_saveall, }, 'image': { 'save': ui_lua_imagesave, }, 'status': ui_status, 'log': { 'info': ui_lua_log_info, 'warning': ui_lua_log_warning, 'error': ui_lua_log_error, } } G.torch = { 'abs': torch_abs, 'add': torch_add, 'sub': torch_sub, 'clamp': torch_clamp, 'lerp': torch_lerp, 'mul': torch_mul, 'div': torch_div, 'size': torch_size, 'new_zeros': torch_new_zeros, 'max': torch_max, 'min': torch_min, 't2f': torch_t2f, 'f2t': torch_f2t, 'cat': torch_cat, } devices.torch_gc() return LUA_output, LUA_gallery def lua_refresh(): global LUA_output, LUA_gallery devices.torch_gc() return LUA_output, LUA_gallery # Functions for Lua def ui_lua_console(text): print(f"Lua: {text}") def ui_lua_output(text): global LUA_output LUA_output += str(text)+'\n' def ui_lua_output_clear(): global LUA_output LUA_output = '' def ui_lua_gallery_add(image): ui_lua_gallery_addc(image, '') def ui_lua_gallery_addc(image, caption): global LUA_gallery LUA_gallery.insert(0, (image, caption)) def ui_lua_gallery_getgif(duration): global LUA_gallery gif = [] for i in LUA_gallery: gif.insert(0, i[0]) gif_path = sd_lua_list2gif(gif, duration) return(gif_path) def ui_lua_gallery_saveall(): path_to_save = os.path.join(opts.outdir_extras_samples, 'lua') if not os.path.exists(path_to_save): try: os.makedirs(path_to_save, exist_ok=True) print('LUA: Creating folder:', path_to_save) except: pass for image in LUA_gallery[::-1]: #gif.insert(0, i[0]) name = images.get_next_sequence_number(path_to_save, '') path_to_image = os.path.join(path_to_save, f"{name}.png") image[0].save(path_to_image) def ui_lua_gallery_clear(): global LUA_gallery LUA_gallery = [] def ui_lua_gallery_del(index): global LUA_gallery # FIXME add code here to match caption del LUA_gallery[index-1] def ui_status(text): shared.state.textinfo = text def ui_lua_log_info(text): #log.info(text) print(f"Lua INFO: {text}") def ui_lua_log_warning(text): #log.warning(text) print(f"Lua WARNING: {text}") def ui_lua_log_error(text): #log.error(text) print(f"Lua ERROR: {text}") def sd_lua_interrogate_clip(image): return(shared.interrogator.interrogate(image)) def sd_lua_interrogate_deepbooru(image): return(deepbooru.model.tag(image)) # Empty latent # IN: width, height # OUT: latent # FIXME, remove this, use torch.new_zeros() instead def sd_lua_empty_latent (w, h): tensor = torch.tensor((), dtype=torch.float32) return tensor.new_zeros((w, h)) def sd_lua_getsamplers(): return([x.name for x in sd_samplers.all_samplers]) def sd_lua_restorefaces(image): return(Image.fromarray(face_restoration.restore_faces(np.asarray(image)))) # IN: # OUT: p def sd_lua_getp(): p = StableDiffusionProcessingTxt2Img( sd_model=shared.sd_model, outpath_samples=shared.opts.outdir_samples or shared.opts.outdir_txt2img_samples, outpath_grids=shared.opts.outdir_grids or shared.opts.outdir_txt2img_grids, prompt='', styles=[], negative_prompt='', seed=-1, subseed=-1, subseed_strength=0, seed_resize_from_h=0, seed_resize_from_w=0, seed_enable_extras=True, sampler_name='Euler a', batch_size=1, n_iter=1, steps=20, cfg_scale=7, width=512, height=512, restore_faces=False, tiling=False, enable_hr=False, denoising_strength=0, hr_scale=0, hr_upscaler=None, hr_second_pass_steps=0, hr_resize_x=0, hr_resize_y=0, override_settings=[], ) return(p) # Conditioning functions ScheduledPromptConditioning = namedtuple("ScheduledPromptConditioning", ["end_at_step", "cond"]) class ComposableScheduledPromptConditioning: def __init__(self, schedules, weight=1.0): self.schedules: List[ScheduledPromptConditioning] = schedules self.weight: float = weight class MulticondLearnedConditioning: def __init__(self, shape, batch): self.shape: tuple = shape # the shape field is needed to send this object to DDIM/PLMS self.batch: List[List[ComposableScheduledPromptConditioning]] = batch def sd_lua_cond(prompt): with devices.autocast(): cond = prompt_parser.get_multicond_learned_conditioning(shared.sd_model, [prompt], 1000) return(cond) def sd_lua_negcond(prompt): with devices.autocast(): cond = prompt_parser.get_learned_conditioning(shared.sd_model, [prompt], 1000) return(cond) def sd_lua_textencode(prompt): # Prompt to learned conditioning with devices.autocast(): conds = shared.sd_model.get_learned_conditioning([prompt]) return(conds) def sd_lua_clip2negcond(conds): # Probably need a better name than clip res = [] cond_schedule = [] with devices.autocast(): for cond in conds: cond_schedule.append(ScheduledPromptConditioning(1000, cond)) # Hardcoded steps, hopefully not breaking too many things res.append(cond_schedule) return(res) def sd_lua_negcond2cond(negcond): res = [] cond_schedule = [] with devices.autocast(): for spcs in negcond: cond_schedule.append([ComposableScheduledPromptConditioning(spcs, 1.0)]) # [[]] return MulticondLearnedConditioning(shape=(1,), batch=cond_schedule) # IN: p, c, uc # OUT: latent def sd_lua_sample(p, c, uc): fix_seed(p) opt_C = 4 opt_f = 8 if ImageRNG is not None: p.rng = ImageRNG((opt_C, p.height // opt_f, p.width // opt_f), [p.seed]) # Fix c and uc so they are of the correc type if c == None: c = '' if uc == None: uc = '' if isinstance(c, str): c = sd_lua_cond(c) if torch.is_tensor(c): c = sd_lua_clip2negcond(c) c = sd_lua_negcond2cond(c) if isinstance(uc, str): uc = sd_lua_negcond(uc) if torch.is_tensor(uc): uc = sd_lua_clip2negcond(uc) with devices.without_autocast() if devices.unet_needs_upcast else devices.autocast(): samples_ddim = p.sample(conditioning=c, unconditional_conditioning=uc, seeds=[p.seed], subseeds=[p.subseed], subseed_strength=p.subseed_strength, prompts=[p.prompt]) devices.torch_gc() return(samples_ddim) # IN: latent # OUT: latent def sd_lua_vae(samples_ddim): x_samples_ddim = [decode_first_stage(shared.sd_model, samples_ddim.to(dtype=devices.dtype_vae))[0].cpu()] try: for x in x_samples_ddim: devices.test_for_nans(x, "vae") except devices.NansException as e: if not shared.cmd_opts.no_half and not shared.cmd_opts.no_half_vae and shared.cmd_opts.rollback_vae: print('\nA tensor with all NaNs was produced in VAE, try converting to bf16.') devices.dtype_vae = torch.bfloat16 vae_file, vae_source = sd_vae.resolve_vae(p.sd_model.sd_model_checkpoint) sd_vae.load_vae(p.sd_model, vae_file, vae_source) x_samples_ddim = [decode_first_stage(p.sd_model, samples_ddim[i:i+1].to(dtype=devices.dtype_vae))[0].cpu() for i in range(samples_ddim.size(0))] for x in x_samples_ddim: devices.test_for_nans(x, "vae") else: raise e x_samples_ddim = torch.stack(x_samples_ddim).float() x_samples_ddim = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0) devices.torch_gc() return(x_samples_ddim) # IN: latent # OUT: image (maybe) def sd_lua_toimage(latent): if len(latent.size()) > 3: # Really ugly kludge latent = latent[0] #for i, x_sample in enumerate(latent): # x_sample = 255. * np.moveaxis(x_sample.cpu().numpy(), 0, 2) # x_sample = x_sample.astype(np.uint8) #image = Image.fromarray(x_sample) T = transforms.ToPILImage() image = T(latent) return image def sd_lua_list2gif(gif, duration): path_to_save = os.path.join(opts.outdir_extras_samples, 'lua') if not os.path.exists(path_to_save): try: os.makedirs(path_to_save, exist_ok=True) print('LUA: Creating folder:', path_to_save) except: pass name = images.get_next_sequence_number(path_to_save, '') path_to_save = os.path.join(path_to_save, f"{name}.gif") gif[0].save(path_to_save, save_all=True, append_images=gif[1:], optimize=False, duration=duration, loop=0) return(path_to_save) # IN: Lua-table with images # OUT: string, name of gif def sd_lua_makegif(table, duration): gif = [x for x in table.values()] gif_path = sd_lua_list2gif(gif, duration) return(gif_path) # IN: image # OUT: string (path to image) def ui_lua_imagesave(image, name): path_to_save = os.path.join(opts.outdir_extras_samples, 'lua') if not os.path.exists(path_to_save): try: os.makedirs(path_to_save, exist_ok=True) print('LUA: Creating folder:', path_to_save) except: pass path_to_save = os.path.join(path_to_save, name) # FIXME check so it doesn't overwrite images? image.save(path_to_save) return(path_to_save) # IN: p # OUT: image def sd_lua_pipeline(p): devices.torch_gc() fix_seed(p) seed = p.seed subseed = p.subseed comments = {} # FIXME remove? ignoring infotext will make things simpler p.all_prompts = [p.prompt] p.all_negative_prompts = [p.negative_prompt] p.all_seeds = [int(seed)] p.all_subseeds = [int(subseed)] def infotext(iteration=0, position_in_batch=0): return create_infotext(p, [p.prompt], [p.seed], [p.subseed], comments, iteration, position_in_batch) infotexts = [] output_images = [] #with torch.no_grad(), p.sd_model.ema_scope(): with torch.no_grad(): prompts = [p.prompt] negative_prompts = [p.negative_prompt] seeds = [p.seed] subseeds = [p.subseed] c = sd_lua_cond(p.prompt) uc = sd_lua_negcond(p.negative_prompt) # Sample samples_ddim = sd_lua_sample(p, c, uc) x_samples_ddim = sd_lua_vae(samples_ddim) del samples_ddim if shared.cmd_opts.lowvram or shared.cmd_opts.medvram: lowvram.send_everything_to_cpu() devices.torch_gc() for i, x_sample in enumerate(x_samples_ddim): x_sample = 255. * np.moveaxis(x_sample.cpu().numpy(), 0, 2) x_sample = x_sample.astype(np.uint8) image = Image.fromarray(x_sample) if opts.samples_save and not p.do_not_save_samples: images.save_image(image, p.outpath_samples, "", seeds[i], prompts[i], opts.samples_format, info=infotext(0, i), p=p) text = infotext(0, i) infotexts.append(text) if opts.enable_pnginfo: image.info["parameters"] = text output_images.append(image) del x_samples_ddim devices.torch_gc() return output_images[0] ############################################################################3 # IN: p or string # OUT: image def sd_lua_process(prompt): if isinstance(prompt, str): p = sd_lua_getp() p.prompt = prompt else: p = prompt processed = process_images(p) p.close() return processed.images[0] ### Torch functions def torch_lerp(v1, v2, weight): return torch.lerp(v1, v2, weight) def torch_abs(v1): return torch.abs(v1) def torch_add(v1, v2): return torch.add(v1, v2) def torch_sub(v1, v2): return torch.sub(v1, v2) def torch_mul(v1, v2): return torch.mul(v1, v2) def torch_div(v1, v2): return torch.div(v1, v2) def torch_clamp(v1, min, max): return torch.clamp(v1, min=min, max=max) def torch_max(v1): return float(torch.max(v1)) def torch_min(v1): return float(torch.min(v1)) def torch_size(v1): return list(v1.size()) def torch_new_zeros(size): tensor = torch.tensor((), dtype=torch.float32) # Lua will return a "table", which for python looks like a dict, so we need to # pick out the values (skipping the keys) and make it into a tuple return tensor.new_zeros(tuple(size.values()), device=torch.device('cuda')) def torch_t2f(tensor): return float(tensor) def torch_f2t(tensor): return float(tensor) def torch_cat(tensor_table, dim): tensors = [x for x in tensor_table.values()] return torch.cat(tuple(tensors), dim) def torch_copy(v): return copy(v) def add_tab(): with gr.Blocks(analytics_enabled=False) as tab: with gr.Row(): with gr.Column(scale=1): # Would use this if the css wasn't broken lua_code = gr.Code(label="Lua", elem_id="lua_code", language=None, show_label=False, lines=30, placeholder="(Lua code)") #lua_code = gr.Textbox(label="Lua", show_label=False, lines=30, placeholder="(Lua code)") with gr.Group(elem_id="sd_webui_lua_results"): with gr.Row(): res_info = gr.HTML() with gr.Row(visible=False): res = gr.Label("") with gr.Column(scale=1): with gr.Row(): lua_input = gr.Textbox(label="Input", lines=1) with gr.Row(): run = gr.Button('Run', variant='primary') reset = gr.Button('Reset') refresh = gr.Button('Refresh') with gr.Row(): gallery = gr.Gallery(label="Gallery").style(preview=True, grid=4) with gr.Row(): results = gr.Textbox(label="Output", show_label=True, lines=10) run.click( show_progress=False, fn=ui.wrap_gradio_gpu_call(lua_run, extra_outputs=['']), _js="submit_sd_webui_lua", inputs=[res, lua_input, lua_code], outputs=[results, gallery, res_info] ) reset.click(lua_reset, show_progress=False, inputs=[], outputs=[results, gallery]) refresh.click(lua_refresh, show_progress=False, inputs=[], outputs=[results, gallery]) with gr.Row(): with gr.Accordion(label='Lua Extras...', open=False): with open(os.path.join(sd_webui_lua_dir, "extras.html"), encoding="utf8") as file: extras = file.read() gr.HTML(extras) return [(tab, "Lua", "lua")] x,y = lua_reset() script_callbacks.on_ui_tabs(add_tab)