app.py

import gradio as gr
import torch
from diffusers import DiffusionPipeline
from diffusers import AutoencoderKL
from config import *
from helpers import *

def device_change(device, config):
    
    config = set_config(config, 'device', device)
    
    return config, config, assemble_code(config)

def model_refiner_change(refiner, config):
    
    config = set_config(config, 'refiner', refiner)
    
    return config, config, assemble_code(config)

def cpu_offload_change(cpu_offload, config):
    
    config = set_config(config, 'cpu_offload', cpu_offload)
    
    return config, config, assemble_code(config)

def models_change(model, scheduler, config):

    config = set_config(config, 'model', model)

    use_safetensors = False
    refiner = "none"
    trigger_token = ""
    
    # no model selected (because this is UI init run)
    if type(model) != list and str(model) != 'None':
        
        use_safetensors = str(models[model]['use_safetensors'])
        model_description = models[model]['description']
        refiner = models[model]['refiner']
        trigger_token = models[model]['trigger_token'] 
        
        # if no scheduler is selected, choose the default one for this model
        if scheduler == None:
            
            scheduler = models[model]['scheduler']

    else:

        model_description = 'Please select a model.'
        
    config["use_safetensors"] = str(use_safetensors)
    config["scheduler"] = str(scheduler)
    config["refiner"] = str(refiner)
    
    # safety_checker_change(in_safety_checker.value, config)
    # requires_safety_checker_change(in_requires_safety_checker.value, config)

    return model_description, refiner, trigger_token, use_safetensors, scheduler, config, config, assemble_code(config)

def data_type_change(data_type, config):

    config = set_config(config, 'data_type', data_type)

    return config, config, assemble_code(config)

def tensorfloat32_change(allow_tensorfloat32, config):  
        
    config = set_config(config, 'allow_tensorfloat32', allow_tensorfloat32)

    return config, config, assemble_code(config)

def inference_steps_change(inference_steps, config):
    
    config = set_config(config, 'inference_steps', inference_steps)

    return config, config, assemble_code(config)

def manual_seed_change(manual_seed, config):
    
    config = set_config(config, 'manual_seed', manual_seed)

    return config, config, assemble_code(config)

def guidance_scale_change(guidance_scale, config):

    config = set_config(config, 'guidance_scale', guidance_scale)

    return config, config, assemble_code(config)

def enable_vae_slicing_change(enable_vae_slicing, config):

    config = set_config(config, 'enable_vae_slicing', enable_vae_slicing)

    return config, config, assemble_code(config)

def enable_vae_tiling_change(enable_vae_tiling, config):

    config = set_config(config, 'enable_vae_tiling', enable_vae_tiling)

    return config, config, assemble_code(config)

def prompt_change(prompt, config):
    
    config = set_config(config, 'prompt', prompt)

    return config, config, assemble_code(config)

def trigger_token_change(trigger_token, config):
    
    config = set_config(config, 'trigger_token', trigger_token)

    return config, config, assemble_code(config)

def negative_prompt_change(negative_prompt, config):
    
    config = set_config(config, 'negative_prompt', negative_prompt)
    
    return config, config, assemble_code(config)

def variant_change(variant, config):
    
    config = set_config(config, 'variant', variant)

    return config, config, assemble_code(config)
    
def safety_checker_change(safety_checker, config):
        
    config = set_config(config, 'safety_checker', safety_checker)

    return config, config, assemble_code(config)

def requires_safety_checker_change(requires_safety_checker, config):

    config = set_config(config, 'requires_safety_checker', requires_safety_checker)

    return config, config, assemble_code(config)

def auto_encoders_change(auto_encoder, config):
    
    if str(auto_encoder) != 'None' and type(auto_encoder) != list:
        
        auto_encoder_description = auto_encoders[auto_encoder]
        
    else:
        auto_encoder_description = ''

    config = set_config(config, 'auto_encoder', auto_encoder)

    return auto_encoder_description, config, config, assemble_code(config)

def schedulers_change(scheduler, config):
    
    if str(scheduler) != 'None' and type(scheduler) != list:
        
        scheduler_description = schedulers[scheduler]
        
    else:
        scheduler_description = 'Please select a scheduler.'
        
    config = set_config(config, 'scheduler', scheduler)

    return scheduler_description, config, config, assemble_code(config)
    
def adapters_textual_inversion_change(adapter_textual_inversion, config):
    
    if str(adapter_textual_inversion) != 'None' and type(adapter_textual_inversion) != list:
        
        adapter_textual_inversion_description = adapters['textual_inversion'][adapter_textual_inversion]['description']
        in_adapters_textual_inversion_token = adapters['textual_inversion'][adapter_textual_inversion]['token']
        
    else:
        adapter_textual_inversion_description = ""
        in_adapters_textual_inversion_token = ""
    
    config = set_config(config, 'adapter_textual_inversion', adapter_textual_inversion)
    
    return adapter_textual_inversion_description, in_adapters_textual_inversion_token, config, config, assemble_code(config)

def textual_inversion_token_change(adapter_textual_inversion_token, config):
    
    config = set_config(config, 'adapter_textual_inversion_token', adapter_textual_inversion_token)

    return config, config, assemble_code(config)

def run_inference(config, config_history, progress=gr.Progress(track_tqdm=True)):
    
    # str_config = str_config.replace("'", '"').replace('None', 'null').replace('False', 'false')
    # config = json.loads(str_config)

    if str(config["model"]) != 'None' and str(config["scheduler"]) != 'None':
        
        progress((1,3), desc="Preparing pipeline initialization...")
        
        torch.backends.cuda.matmul.allow_tf32 = get_bool(config["allow_tensorfloat32"]) # Use TensorFloat-32 as of https://huggingface.co/docs/diffusers/main/en/optimization/fp16 faster, but slightly less accurate computations
        
        progress((2,3), desc="Initializing pipeline...")
        
        # INIT PIPELINE
        pipeline = DiffusionPipeline.from_pretrained(
            config["model"], 
            use_safetensors = get_bool(config["use_safetensors"]), 
            torch_dtype = get_data_type(config["data_type"]), 
            variant = get_variant(config["variant"])).to(config["device"])
        
        if str(config["cpu_offload"]).lower() != 'false':
            pipeline.enable_model_cpu_offload()

        # AUTO ENCODER
        if str(config["auto_encoder"]).lower() != 'none': 
            pipeline.vae = AutoencoderKL.from_pretrained(config["auto_encoder"], torch_dtype=get_data_type(config["data_type"])).to(config["device"])

        if str(config["enable_vae_slicing"]).lower() != 'false': pipeline.enable_vae_slicing()
        if str(config["enable_vae_tiling"]).lower() != 'false': pipeline.enable_vae_tiling()
                
        # INIT REFINER
        if config['refiner'].lower() != 'none':
            refiner = DiffusionPipeline.from_pretrained(
                    config['refiner'],
                    text_encoder_2=pipeline.text_encoder_2,
                    vae=pipeline.vae,
                    torch_dtype=get_data_type(config["data_type"]),
                    use_safetensors=get_bool(config["use_safetensors"]), 
                    variant = get_variant(config["variant"])).to(config["device"])
            
            if str(config["cpu_offload"]).lower() != 'false':
                refiner.enable_model_cpu_offload()

            if str(config["enable_vae_slicing"]).lower() != 'false': refiner.enable_vae_slicing()
            if str(config["enable_vae_tiling"]).lower() != 'false': refiner.enable_vae_tiling()
            
        # SAFETY CHECKER
        if str(config["safety_checker"]).lower() == 'false': pipeline.safety_checker = None 
        pipeline.requires_safety_checker = get_bool(config["requires_safety_checker"])
            
        # SCHEDULER/SOLVER
        pipeline.scheduler = get_scheduler(config["scheduler"], pipeline.scheduler.config)
        
        # MANUAL SEED/GENERATOR
        if config["manual_seed"] is None or config["manual_seed"] == '' or int(config["manual_seed"]) < 0:
            generator = None
        else:
            generator = torch.manual_seed(int(config["manual_seed"])) 
        
        # ADAPTERS
        # TEXTUAL INVERSION
        if str(config["adapter_textual_inversion"]).lower() != 'none':
            pipeline.load_textual_inversion(config["adapter_textual_inversion"], token=config["adapter_textual_inversion_token"])
        
        progress((3,3), desc="Creating the result...")

        prompt = config["prompt"] + config["trigger_token"] + config["adapter_textual_inversion_token"]
        
        image = pipeline(
            prompt = prompt,
            negative_prompt = config["negative_prompt"],
            generator = generator,
            num_inference_steps = int(config["inference_steps"]), 
            guidance_scale = float(config["guidance_scale"])).images

        if config['refiner'].lower() != 'none':
            image = refiner(
                prompt = prompt,
                num_inference_steps = int(config["inference_steps"]), 
                image=image,
            ).images
            
        config_history.append(config.copy())

        return image[0], dict_list_to_markdown_table(config_history), config_history
    
    else:
    
        return "Please select a model AND a scheduler.", None, config_history

appConfig = load_app_config()
models = appConfig.get("models", {})
schedulers = appConfig.get("schedulers", {})
devices =  appConfig.get("devices", [])
auto_encoders = appConfig.get("auto_encoders", [])
adapters = appConfig.get("adapters", [])

js = '''function js(){
        window.set_cookie = function(key, value, config){
            document.cookie = key+'='+value+'; Path=/; SameSite=Strict';
            return [value, config]
        }
        window.set_model_cookie = function(model, config){
            document.cookie = 'model='+ model+'; Path=/; SameSite=Strict';
            // some things I just don't understand, this is one of them
            return [model, null, null, null, null, model, config, null]
        }
        window.set_adapter_textual_inversion_cookie = function(adapter_textual_inversion, config){
            document.cookie = 'adapter_textual_inversion='+ adapter_textual_inversion+'; Path=/; SameSite=Strict';
            // some things I just don't understand, this is one of them
            return [adapter_textual_inversion, null, adapter_textual_inversion, config, null]
        }
        window.set_cookie_2 = function(key, value, config){
            document.cookie = key+'='+value+'; Path=/; SameSite=Strict';
            // some things I just don't understand, this is one of them
            return [value, null, config, null]
        }
    }
    '''

# interface
with gr.Blocks(analytics_enabled=False) as demo:
    
    config = gr.State(value=get_initial_config())
    config_history = gr.State(value=[])

    gr.Markdown('''## Text-2-Image Playground
                <small>by <a target="_blank" href="https://www.linkedin.com/in/nickyreinert/">Nicky Reinert</a> | 
                home base: https://huggingface.co/spaces/n42/pictero
                </small>''')
    gr.Markdown("### Device specific settings")
    with gr.Row():
        in_devices = gr.Dropdown(label="Device:", value=config.value["device"], choices=devices, filterable=True, multiselect=False, allow_custom_value=True, info="(you may add a custom device address at any time)")
        in_data_type = gr.Radio(label="Data Type:", value=config.value["data_type"], choices=["bfloat16", "float16", "float32"], info="`bfloat16` is not supported on MPS devices right now; `float16` may also not be supported on all devices, Half-precision weights, will save GPU memory, see https://huggingface.co/docs/diffusers/main/en/optimization/fp16")
        in_allow_tensorfloat32 = gr.Radio(label="Allow TensorFloat32:", value=config.value["allow_tensorfloat32"], choices=["True", "False"], info="is not supported on MPS devices right now; use TensorFloat-32 is faster, but results in slightly less accurate computations, see https://huggingface.co/docs/diffusers/main/en/optimization/fp16 ")
        in_variant = gr.Radio(label="Variant:", value=config.value["variant"], choices=["fp16", None], info="Use half-precision weights will save GPU memory, not all models support that, see https://huggingface.co/docs/diffusers/main/en/optimization/fp16 ")
        
    gr.Markdown("### Model specific settings")
    with gr.Row():
        in_models = gr.Dropdown(choices=list(models.keys()), label="Model")
        out_model_description = gr.Textbox(value="", label="Description")
    with gr.Row():
        with gr.Column(scale=1):
            in_trigger_token = gr.Textbox(value=config.value["trigger_token"], label="Trigger Token", info="will be added to your prompt to `activate` a fine tuned model")
            in_use_safetensors = gr.Radio(label="Use safe tensors:", choices=["True", "False"], interactive=False)
            in_model_refiner = gr.Dropdown(value=config.value["refiner"], choices=["none"], label="Refiner", allow_custom_value=True, multiselect=False)
        with gr.Column(scale=1):
            in_cpu_offload = gr.Radio(label="CPU Offload:", value=config.value["cpu_offload"], choices=["True", "False"], info="This may increase performance, as it offloads computations from the GPU to the CPU. But this can also lead to slower executions and lower effectiveness. Compare running time and outputs before making sure, that this setting will help you")
            in_safety_checker = gr.Radio(label="Enable safety checker:", value=config.value["safety_checker"], choices=["True", "False"])
            in_requires_safety_checker = gr.Radio(label="Requires safety checker:", value=config.value["requires_safety_checker"], choices=["True", "False"])

    gr.Markdown("### Scheduler")
    with gr.Row():
        in_schedulers = gr.Dropdown(value="", choices=list(schedulers.keys()), label="Scheduler/Solver", info="schedulers employ various strategies for noise control, the scheduler controls parameter adaption between each inference step, depending on the right scheduler for your model, it may only take 10 or 20 steps to achieve very good results, see https://huggingface.co/docs/diffusers/using-diffusers/loading#schedulers" )
        out_scheduler_description = gr.Textbox(value="", label="Description")
    
    # gr.Markdown("### Adapters")
    # with gr.Row():
    #     gr.Markdown('Choose an adapter.')
        
    gr.Markdown("### Auto Encoder")
    with gr.Row():
        gr.Markdown("**VAE** stands for Variational Auto Encoders. An 'autoencoder' is an artificial neural network that is able to encode input data and decode to output data to bascially recreate the input. The VAE whereas adds a couple of additional layers of complexity to create new and unique output.")
    with gr.Row():
        with gr.Column():
            in_auto_encoders = gr.Dropdown(value="None", choices=list(auto_encoders.keys()), label="Auto encoder", info="leave empty to not add an auto encoder")
            out_auto_encoder_description = gr.Textbox(value="", label="Description")
        with gr.Column():
            in_enable_vae_slicing = gr.Radio(label="Enable VAE slicing:", value=config.value["enable_vae_slicing"], choices=["True", "False"], info="decoding the batches of latents one image at a time, which may reduce memory usage, see https://huggingface.co/docs/diffusers/main/en/optimization/memory")
            in_enable_vae_tiling= gr.Radio(label="Enable VAE tiling:", value=config.value["enable_vae_tiling"], choices=["True", "False"], info="splitting the image into overlapping tiles, decoding the tiles, and then blending the outputs together to compose the final image, see https://huggingface.co/docs/diffusers/main/en/optimization/memory")
        
    gr.Markdown("### Adapters")
    with gr.Row():
        gr.Markdown('''Adapters allow you to apply finetuned weights to your base model. They come in many flavors depending on how they were trained. See see https://huggingface.co/docs/diffusers/using-diffusers/loading_adapters''')
    with gr.Row():
        gr.Markdown('#### Textual Inversion Adapters')
    with gr.Row():
        gr.Markdown('(a technique that enables a model like Stable Diffusion to learn a new concept from just a few sample images)')
    with gr.Row():
        in_adapters_textual_inversion = gr.Dropdown(value="None", choices=list(adapters['textual_inversion'].keys()), label="Textual Inversion Adapter", info="leave empty to not use an adapter")
        in_adapters_textual_inversion_token = gr.Textbox(value="None", label="Token", info="required to activate the token, will be added to your prompt")
        out_adapters_textual_inversion_description = gr.Textbox(value="", label="Description")
        
    gr.Markdown("### Inference settings")
    with gr.Row():
        in_prompt = gr.TextArea(label="Prompt", value=config.value["prompt"])
        in_negative_prompt = gr.TextArea(label="Negative prompt", value=config.value["negative_prompt"])
    with gr.Row():
        in_inference_steps = gr.Number(label="Inference steps", value=config.value["inference_steps"], info="Each step improves the final result but also results in higher computation")
        in_manual_seed = gr.Number(label="Manual seed", value=config.value["manual_seed"], info="Set this to -1 or leave it empty to randomly generate an image. A fixed value will result in a similar image for every run")
        in_guidance_scale = gr.Slider(minimum=0, maximum=100, step=0.1, label="Guidance Scale", value=config.value["guidance_scale"], info="A low guidance scale leads to a faster inference time, with the drawback that negative prompts don’t have any effect on the denoising process.")
    
    gr.Markdown("### Output")
    with gr.Row():
        btn_start_pipeline = gr.Button(value="Run", variant="primary")
        btn_stop_pipeline = gr.Button(value="Stop", variant="stop")
    with gr.Row():
        out_image = gr.Image()
        out_code = gr.Code(assemble_code(config.value), label="Code")
    with gr.Row():
        # out_config = gr.Code(value=str(config.value), label="Current config")
        out_config = gr.JSON(value=config.value, label="Current config")
    with gr.Row():
        out_config_history = gr.Markdown(dict_list_to_markdown_table(config_history.value))
    
    # `SPECIAL` CHANGE LISTENERS
    in_models.change(models_change, inputs=[in_models, in_schedulers, config], outputs=[out_model_description, in_model_refiner, in_trigger_token, in_use_safetensors, in_schedulers, config, out_config, out_code], js="(model, config) => set_model_cookie(model, config)")
    in_schedulers.change(schedulers_change, inputs=[in_schedulers, config], outputs=[out_scheduler_description, config, out_config, out_code], js="(value, config) => set_cookie_2('scheduler', value, config)")
    in_auto_encoders.change(auto_encoders_change, inputs=[in_auto_encoders, config], outputs=[out_auto_encoder_description, config, out_config, out_code], js="(value, config) => set_cookie_2('auto_encoder', value, config)")
    in_adapters_textual_inversion.change(adapters_textual_inversion_change, inputs=[in_adapters_textual_inversion, config], outputs=[out_adapters_textual_inversion_description, in_adapters_textual_inversion_token, config, out_config, out_code], js="(adapter_textual_inversion, config) => set_adapter_textual_inversion_cookie(adapter_textual_inversion, config)")
    
    # `GENERIC` CHANGE LISTENERS, SAME INPUT, SAME OUTPUT STRUCTURE
    in_devices.change(fn=device_change, inputs=[in_devices, config], outputs=[config, out_config, out_code], js="(value, config) => set_cookie('device', value, config)")
    in_data_type.change(data_type_change, inputs=[in_data_type, config], outputs=[config, out_config, out_code], js="(value, config) => set_cookie('data_type', value, config)")
    in_allow_tensorfloat32.change(tensorfloat32_change, inputs=[in_allow_tensorfloat32, config], outputs=[config, out_config, out_code], js="(value, config) => set_cookie('allow_tensorfloat32', value, config)")
    in_variant.change(variant_change, inputs=[in_variant, config], outputs=[config, out_config, out_code], js="(value, config) => set_cookie('variant', value, config)")
    in_model_refiner.change(model_refiner_change, inputs=[in_model_refiner, config], outputs=[config, out_config, out_code], js="(value, config) => set_cookie('model_refiner', value, config)")
    in_cpu_offload.change(cpu_offload_change, inputs=[in_cpu_offload, config], outputs=[config, out_config, out_code], js="(value, config) => set_cookie('cpu_offload', value, config)")
    in_safety_checker.change(safety_checker_change, inputs=[in_safety_checker, config], outputs=[config, out_config, out_code], js="(value, config) => set_cookie('safety_checker', value, config)")
    in_requires_safety_checker.change(requires_safety_checker_change, inputs=[in_requires_safety_checker, config], outputs=[config, out_config, out_code], js="(value, config) => set_cookie('requires_safety_checker', value, config)")
    in_inference_steps.change(inference_steps_change, inputs=[in_inference_steps, config], outputs=[config, out_config, out_code], js="(value, config) => set_cookie('inference_steps', value, config)")
    in_manual_seed.change(manual_seed_change, inputs=[in_manual_seed, config], outputs=[config, out_config, out_code], js="(value, config) => set_cookie('manual_seed', value, config)")
    in_guidance_scale.change(guidance_scale_change, inputs=[in_guidance_scale, config], outputs=[config, out_config, out_code], js="(value, config) => set_cookie('guidance_scale', value, config)")
    in_enable_vae_slicing.change(enable_vae_slicing_change, inputs=[in_enable_vae_slicing, config], outputs=[config, out_config, out_code], js="(value, config) => set_cookie('enable_vae_slicing', value, config)")
    in_enable_vae_tiling.change(enable_vae_tiling_change, inputs=[in_enable_vae_tiling, config], outputs=[config, out_config, out_code], js="(value, config) => set_cookie('enable_vae_tiling', value, config)")
    in_adapters_textual_inversion_token.change(textual_inversion_token_change, inputs=[in_adapters_textual_inversion_token, config], outputs=[config, out_config, out_code])
    in_prompt.change(prompt_change, inputs=[in_prompt, config], outputs=[config, out_config, out_code], js="(value, config) => set_cookie('prompt', value, config)")
    in_trigger_token.change(trigger_token_change, inputs=[in_trigger_token, config], outputs=[config, out_config, out_code], js="(value, config) => set_cookie('trigger_token', value, config)")
    in_negative_prompt.change(negative_prompt_change, inputs=[in_negative_prompt, config], outputs=[config, out_config, out_code], js="(value, config) => set_cookie('negative_prompt', value, config)")
    
    ev_run_inference = btn_start_pipeline.click(run_inference, inputs=[config, config_history], outputs=[out_image, out_config_history, config_history])
    btn_stop_pipeline.click(fn=None, inputs=None, outputs=None, cancels=[ev_run_inference])

    # send current respect initial config to init_config to populate parameters to all relevant input fields
    # if GET parameter is set, it will overwrite initial config parameters
    demo.load(fn=get_config_from_url, js=js,
        inputs=[config], 
        outputs=[
            in_models,
            in_devices,
            in_cpu_offload,
            in_use_safetensors,
            in_data_type,
            in_model_refiner,
            in_variant,
            in_safety_checker,
            in_requires_safety_checker,
            in_auto_encoders,
            in_enable_vae_slicing,
            in_enable_vae_tiling,
            in_schedulers,
            in_prompt,
            in_trigger_token,
            in_negative_prompt,
            in_inference_steps,
            in_manual_seed,
            in_guidance_scale,
            in_adapters_textual_inversion
            ])

demo.launch(show_error=True)