diff --git a/apps/stable_diffusion/scripts/train_lora_word.py b/apps/stable_diffusion/scripts/train_lora_word.py index b03edf11..6b37324f 100644 --- a/apps/stable_diffusion/scripts/train_lora_word.py +++ b/apps/stable_diffusion/scripts/train_lora_word.py @@ -223,7 +223,8 @@ def lora_train( if not hf_model_id: return ( None, - "Please provide either custom model or huggingface model ID, both must not be empty", + "Please provide either custom model or huggingface model ID, both must not be " + "empty.", ) args.hf_model_id = hf_model_id elif ".ckpt" in custom_model or ".safetensors" in custom_model: diff --git a/apps/stable_diffusion/scripts/upscaler.py b/apps/stable_diffusion/scripts/upscaler.py index baf433ae..3ea57033 100644 --- a/apps/stable_diffusion/scripts/upscaler.py +++ b/apps/stable_diffusion/scripts/upscaler.py @@ -21,7 +21,7 @@ if __name__ == "__main__": print("Flag --img_path is required.") exit() - # When the models get uploaded, it should be default to False. + # When the models get uploaded, it should be defaulted to False. args.import_mlir = True cpu_scheduling = not args.scheduler.startswith("Shark") diff --git a/apps/stable_diffusion/src/pipelines/pipeline_shark_stable_diffusion_img2img.py b/apps/stable_diffusion/src/pipelines/pipeline_shark_stable_diffusion_img2img.py index 24ad167c..64e63eb6 100644 --- a/apps/stable_diffusion/src/pipelines/pipeline_shark_stable_diffusion_img2img.py +++ b/apps/stable_diffusion/src/pipelines/pipeline_shark_stable_diffusion_img2img.py @@ -15,6 +15,11 @@ from diffusers import ( EulerAncestralDiscreteScheduler, DPMSolverMultistepScheduler, DEISMultistepScheduler, + DPMSolverSinglestepScheduler, + KDPM2AncestralDiscreteScheduler, + HeunDiscreteScheduler, + DDPMScheduler, + KDPM2DiscreteScheduler, ) from apps.stable_diffusion.src.schedulers import SharkEulerDiscreteScheduler from apps.stable_diffusion.src.pipelines.pipeline_shark_stable_diffusion_utils import ( @@ -38,6 +43,11 @@ class Image2ImagePipeline(StableDiffusionPipeline): DPMSolverMultistepScheduler, SharkEulerDiscreteScheduler, DEISMultistepScheduler, + DPMSolverSinglestepScheduler, + KDPM2AncestralDiscreteScheduler, + HeunDiscreteScheduler, + DDPMScheduler, + KDPM2DiscreteScheduler, ], sd_model: SharkifyStableDiffusionModel, import_mlir: bool, diff --git a/apps/stable_diffusion/src/pipelines/pipeline_shark_stable_diffusion_inpaint.py b/apps/stable_diffusion/src/pipelines/pipeline_shark_stable_diffusion_inpaint.py index 515a7bd5..b43c643b 100644 --- a/apps/stable_diffusion/src/pipelines/pipeline_shark_stable_diffusion_inpaint.py +++ b/apps/stable_diffusion/src/pipelines/pipeline_shark_stable_diffusion_inpaint.py @@ -14,6 +14,11 @@ from diffusers import ( EulerAncestralDiscreteScheduler, DPMSolverMultistepScheduler, DEISMultistepScheduler, + DPMSolverSinglestepScheduler, + KDPM2AncestralDiscreteScheduler, + HeunDiscreteScheduler, + DDPMScheduler, + KDPM2DiscreteScheduler, ) from apps.stable_diffusion.src.schedulers import SharkEulerDiscreteScheduler from apps.stable_diffusion.src.pipelines.pipeline_shark_stable_diffusion_utils import ( @@ -37,6 +42,11 @@ class InpaintPipeline(StableDiffusionPipeline): DPMSolverMultistepScheduler, SharkEulerDiscreteScheduler, DEISMultistepScheduler, + DPMSolverSinglestepScheduler, + KDPM2AncestralDiscreteScheduler, + HeunDiscreteScheduler, + DDPMScheduler, + KDPM2DiscreteScheduler, ], sd_model: SharkifyStableDiffusionModel, import_mlir: bool, diff --git a/apps/stable_diffusion/src/pipelines/pipeline_shark_stable_diffusion_outpaint.py b/apps/stable_diffusion/src/pipelines/pipeline_shark_stable_diffusion_outpaint.py index d782633d..982bec6d 100644 --- a/apps/stable_diffusion/src/pipelines/pipeline_shark_stable_diffusion_outpaint.py +++ b/apps/stable_diffusion/src/pipelines/pipeline_shark_stable_diffusion_outpaint.py @@ -14,6 +14,11 @@ from diffusers import ( EulerAncestralDiscreteScheduler, DPMSolverMultistepScheduler, DEISMultistepScheduler, + DPMSolverSinglestepScheduler, + KDPM2AncestralDiscreteScheduler, + HeunDiscreteScheduler, + DDPMScheduler, + KDPM2DiscreteScheduler, ) from apps.stable_diffusion.src.schedulers import SharkEulerDiscreteScheduler from apps.stable_diffusion.src.pipelines.pipeline_shark_stable_diffusion_utils import ( @@ -38,6 +43,11 @@ class OutpaintPipeline(StableDiffusionPipeline): DPMSolverMultistepScheduler, SharkEulerDiscreteScheduler, DEISMultistepScheduler, + DPMSolverSinglestepScheduler, + KDPM2AncestralDiscreteScheduler, + HeunDiscreteScheduler, + DDPMScheduler, + KDPM2DiscreteScheduler, ], sd_model: SharkifyStableDiffusionModel, import_mlir: bool, diff --git a/apps/stable_diffusion/src/pipelines/pipeline_shark_stable_diffusion_stencil.py b/apps/stable_diffusion/src/pipelines/pipeline_shark_stable_diffusion_stencil.py index e7fa415f..6b4eb0a1 100644 --- a/apps/stable_diffusion/src/pipelines/pipeline_shark_stable_diffusion_stencil.py +++ b/apps/stable_diffusion/src/pipelines/pipeline_shark_stable_diffusion_stencil.py @@ -14,6 +14,12 @@ from diffusers import ( EulerDiscreteScheduler, EulerAncestralDiscreteScheduler, DPMSolverMultistepScheduler, + DEISMultistepScheduler, + DPMSolverSinglestepScheduler, + KDPM2AncestralDiscreteScheduler, + HeunDiscreteScheduler, + DDPMScheduler, + KDPM2DiscreteScheduler, ) from apps.stable_diffusion.src.schedulers import SharkEulerDiscreteScheduler from apps.stable_diffusion.src.pipelines.pipeline_shark_stable_diffusion_utils import ( @@ -38,6 +44,12 @@ class StencilPipeline(StableDiffusionPipeline): EulerAncestralDiscreteScheduler, DPMSolverMultistepScheduler, SharkEulerDiscreteScheduler, + DEISMultistepScheduler, + DPMSolverSinglestepScheduler, + KDPM2AncestralDiscreteScheduler, + HeunDiscreteScheduler, + DDPMScheduler, + KDPM2DiscreteScheduler, ], sd_model: SharkifyStableDiffusionModel, import_mlir: bool, diff --git a/apps/stable_diffusion/src/schedulers/sd_schedulers.py b/apps/stable_diffusion/src/schedulers/sd_schedulers.py index 0e7ab95e..9ee91c0f 100644 --- a/apps/stable_diffusion/src/schedulers/sd_schedulers.py +++ b/apps/stable_diffusion/src/schedulers/sd_schedulers.py @@ -8,6 +8,9 @@ from diffusers import ( EulerDiscreteScheduler, EulerAncestralDiscreteScheduler, DEISMultistepScheduler, + DPMSolverSinglestepScheduler, + KDPM2AncestralDiscreteScheduler, + HeunDiscreteScheduler, ) from apps.stable_diffusion.src.schedulers.shark_eulerdiscrete import ( SharkEulerDiscreteScheduler, @@ -62,5 +65,21 @@ def get_schedulers(model_id): model_id, subfolder="scheduler", ) + schedulers[ + "DPMSolverSinglestep" + ] = DPMSolverSinglestepScheduler.from_pretrained( + model_id, + subfolder="scheduler", + ) + schedulers[ + "KDPM2AncestralDiscrete" + ] = KDPM2AncestralDiscreteScheduler.from_pretrained( + model_id, + subfolder="scheduler", + ) + schedulers["HeunDiscrete"] = HeunDiscreteScheduler.from_pretrained( + model_id, + subfolder="scheduler", + ) schedulers["SharkEulerDiscrete"].compile() return schedulers diff --git a/apps/stable_diffusion/src/utils/resources/prompts.json b/apps/stable_diffusion/src/utils/resources/prompts.json index 4c8370db..83871784 100644 --- a/apps/stable_diffusion/src/utils/resources/prompts.json +++ b/apps/stable_diffusion/src/utils/resources/prompts.json @@ -5,4 +5,7 @@ ["A digital Illustration of the Babel tower, 4k, detailed, trending in artstation, fantasy vivid colors"], ["Cluttered house in the woods, anime, oil painting, high resolution, cottagecore, ghibli inspired, 4k"], ["A beautiful mansion beside a waterfall in the woods, by josef thoma, matte painting, trending on artstation HQ"], -["portrait photo of a asia old warrior chief, tribal panther make up, blue on red, side profile, looking away, serious eyes"]] +["portrait photo of a asia old warrior chief, tribal panther make up, blue on red, side profile, looking away, serious eyes"], +["A photo of a beach, sunset, calm, beautiful landscape, waves, water"], +["(a large body of water with snowy mountains in the background), (fog, foggy, rolling fog), (clouds, cloudy, rolling clouds), dramatic sky and landscape, extraordinary landscape, (beautiful snow capped mountain background), (forest, dirt path)"], +["a photo taken of the front of a super-car drifting on a road near mountains at high speeds with smokes coming off the tires, front angle, front point of view, trees in the mountains of the background, ((sharp focus))"]] diff --git a/apps/stable_diffusion/src/utils/sd_annotation.py b/apps/stable_diffusion/src/utils/sd_annotation.py index e1fcc557..5c6bea95 100644 --- a/apps/stable_diffusion/src/utils/sd_annotation.py +++ b/apps/stable_diffusion/src/utils/sd_annotation.py @@ -131,11 +131,32 @@ def load_lower_configs(base_model_id=None): "v1_4", "v1_5", ]: - config_name = f"{args.annotation_model}_{version}_{args.max_length}_{args.precision}_{device}_{spec}_{args.width}x{args.height}.json" + config_name = ( + f"{args.annotation_model}_" + f"{version}_" + f"{args.max_length}_" + f"{args.precision}_" + f"{device}_" + f"{spec}_" + f"{args.width}x{args.height}.json" + ) elif spec in ["rdna2"] and version in ["v2_1", "v2_1base", "v1_4"]: - config_name = f"{args.annotation_model}_{version}_{args.precision}_{device}_{spec}_{args.width}x{args.height}.json" + config_name = ( + f"{args.annotation_model}_" + f"{version}_" + f"{args.precision}_" + f"{device}_" + f"{spec}_" + f"{args.width}x{args.height}.json" + ) else: - config_name = f"{args.annotation_model}_{version}_{args.precision}_{device}_{spec}.json" + config_name = ( + f"{args.annotation_model}_" + f"{version}_" + f"{args.precision}_" + f"{device}_" + f"{spec}.json" + ) full_gs_url = config_bucket + config_name lowering_config_dir = os.path.join(WORKDIR, "configs", config_name) @@ -180,9 +201,22 @@ def dump_after_mlir(input_mlir, use_winograd): device, device_spec_args = get_device_args() if use_winograd: - preprocess_flag = "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-flow-detach-elementwise-from-named-ops,iree-flow-convert-1x1-filter-conv2d-to-matmul,iree-preprocessing-convert-conv2d-to-img2col,iree-preprocessing-pad-linalg-ops{pad-size=32},iree-linalg-ext-convert-conv2d-to-winograd))" + preprocess_flag = ( + "--iree-preprocessing-pass-pipeline=builtin.module" + "(func.func(iree-flow-detach-elementwise-from-named-ops," + "iree-flow-convert-1x1-filter-conv2d-to-matmul," + "iree-preprocessing-convert-conv2d-to-img2col," + "iree-preprocessing-pad-linalg-ops{pad-size=32}," + "iree-linalg-ext-convert-conv2d-to-winograd))" + ) else: - preprocess_flag = "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-flow-detach-elementwise-from-named-ops,iree-flow-convert-1x1-filter-conv2d-to-matmul,iree-preprocessing-convert-conv2d-to-img2col,iree-preprocessing-pad-linalg-ops{pad-size=32}))" + preprocess_flag = ( + "--iree-preprocessing-pass-pipeline=builtin.module" + "(func.func(iree-flow-detach-elementwise-from-named-ops," + "iree-flow-convert-1x1-filter-conv2d-to-matmul," + "iree-preprocessing-convert-conv2d-to-img2col," + "iree-preprocessing-pad-linalg-ops{pad-size=32}))" + ) dump_module = ireec.compile_str( input_mlir, diff --git a/apps/stable_diffusion/src/utils/stable_args.py b/apps/stable_diffusion/src/utils/stable_args.py index 0a905535..c712ee36 100644 --- a/apps/stable_diffusion/src/utils/stable_args.py +++ b/apps/stable_diffusion/src/utils/stable_args.py @@ -19,48 +19,55 @@ p = argparse.ArgumentParser( ) ############################################################################## -### Stable Diffusion Params +# Stable Diffusion Params ############################################################################## p.add_argument( "-a", "--app", default="txt2img", - help="which app to use, one of: txt2img, img2img, outpaint, inpaint", + help="Which app to use, one of: txt2img, img2img, outpaint, inpaint.", ) p.add_argument( "-p", "--prompts", nargs="+", - default=["cyberpunk forest by Salvador Dali"], - help="text of which images to be generated.", + default=[ + "a photo taken of the front of a super-car drifting on a road near mountains at " + "high speeds with smokes coming off the tires, front angle, front point of view, " + "trees in the mountains of the background, ((sharp focus))" + ], + help="Text of which images to be generated.", ) p.add_argument( "--negative_prompts", nargs="+", - default=["trees, green"], - help="text you don't want to see in the generated image.", + default=[ + "watermark, signature, logo, text, lowres, ((monochrome, grayscale)), blurry, " + "ugly, blur, oversaturated, cropped" + ], + help="Text you don't want to see in the generated image.", ) p.add_argument( "--img_path", type=str, - help="Path to the image input for img2img/inpainting", + help="Path to the image input for img2img/inpainting.", ) p.add_argument( "--steps", type=int, - default=50, - help="the no. of steps to do the sampling.", + default=32, + help="The number of steps to do the sampling.", ) p.add_argument( "--seed", type=int, default=-1, - help="the seed to use. -1 for a random one.", + help="The seed to use. -1 for a random one.", ) p.add_argument( @@ -68,7 +75,7 @@ p.add_argument( type=int, default=1, choices=range(1, 4), - help="the number of inferences to be made in a single `batch_count`.", + help="The number of inferences to be made in a single `batch_count`.", ) p.add_argument( @@ -76,7 +83,7 @@ p.add_argument( type=int, default=512, choices=range(128, 769, 8), - help="the height of the output image.", + help="The height of the output image.", ) p.add_argument( @@ -84,84 +91,85 @@ p.add_argument( type=int, default=512, choices=range(128, 769, 8), - help="the width of the output image.", + help="The width of the output image.", ) p.add_argument( "--guidance_scale", type=float, - default=7.5, - help="the value to be used for guidance scaling.", + default=3.5, + help="The value to be used for guidance scaling.", ) p.add_argument( "--noise_level", type=int, default=20, - help="the value to be used for noise level of upscaler.", + help="The value to be used for noise level of upscaler.", ) p.add_argument( "--max_length", type=int, default=64, - help="max length of the tokenizer output, options are 64 and 77.", + help="Max length of the tokenizer output, options are 64 and 77.", ) p.add_argument( "--max_embeddings_multiples", type=int, default=5, - help="The max multiple length of prompt embeddings compared to the max output length of text encoder.", + help="The max multiple length of prompt embeddings compared to the max output " + "length of text encoder.", ) p.add_argument( "--strength", type=float, default=0.8, - help="the strength of change applied on the given input image for img2img", + help="The strength of change applied on the given input image for img2img.", ) ############################################################################## -### Stable Diffusion Training Params +# Stable Diffusion Training Params ############################################################################## p.add_argument( "--lora_save_dir", type=str, default="models/lora/", - help="Directory to save the lora fine tuned model", + help="Directory to save the lora fine tuned model.", ) p.add_argument( "--training_images_dir", type=str, default="models/lora/training_images/", - help="Directory containing images that are an example of the prompt", + help="Directory containing images that are an example of the prompt.", ) p.add_argument( "--training_steps", type=int, default=2000, - help="The no. of steps to train", + help="The number of steps to train.", ) ############################################################################## -### Inpainting and Outpainting Params +# Inpainting and Outpainting Params ############################################################################## p.add_argument( "--mask_path", type=str, - help="Path to the mask image input for inpainting", + help="Path to the mask image input for inpainting.", ) p.add_argument( "--inpaint_full_res", default=False, action=argparse.BooleanOptionalAction, - help="If inpaint only masked area or whole picture", + help="If inpaint only masked area or whole picture.", ) p.add_argument( @@ -169,7 +177,7 @@ p.add_argument( type=int, default=32, choices=range(0, 257, 4), - help="Number of pixels for only masked padding", + help="Number of pixels for only masked padding.", ) p.add_argument( @@ -177,7 +185,7 @@ p.add_argument( type=int, default=128, choices=range(8, 257, 8), - help="Number of expended pixels for one direction for outpainting", + help="Number of expended pixels for one direction for outpainting.", ) p.add_argument( @@ -185,89 +193,91 @@ p.add_argument( type=int, default=8, choices=range(0, 65), - help="Number of blur pixels for outpainting", + help="Number of blur pixels for outpainting.", ) p.add_argument( "--left", default=False, action=argparse.BooleanOptionalAction, - help="If expend left for outpainting", + help="If expend left for outpainting.", ) p.add_argument( "--right", default=False, action=argparse.BooleanOptionalAction, - help="If expend right for outpainting", + help="If expend right for outpainting.", ) p.add_argument( "--top", default=False, action=argparse.BooleanOptionalAction, - help="If expend top for outpainting", + help="If expend top for outpainting.", ) p.add_argument( "--bottom", default=False, action=argparse.BooleanOptionalAction, - help="If expend bottom for outpainting", + help="If expend bottom for outpainting.", ) p.add_argument( "--noise_q", type=float, default=1.0, - help="Fall-off exponent for outpainting (lower=higher detail) (min=0.0, max=4.0)", + help="Fall-off exponent for outpainting (lower=higher detail) (min=0.0, max=4.0).", ) p.add_argument( "--color_variation", type=float, default=0.05, - help="Color variation for outpainting (min=0.0, max=1.0)", + help="Color variation for outpainting (min=0.0, max=1.0).", ) ############################################################################## -### Model Config and Usage Params +# Model Config and Usage Params ############################################################################## p.add_argument( - "--device", type=str, default="vulkan", help="device to run the model." + "--device", type=str, default="vulkan", help="Device to run the model." ) p.add_argument( - "--precision", type=str, default="fp16", help="precision to run the model." + "--precision", type=str, default="fp16", help="Precision to run the model." ) p.add_argument( "--import_mlir", default=False, action=argparse.BooleanOptionalAction, - help="imports the model from torch module to shark_module otherwise downloads the model from shark_tank.", + help="Imports the model from torch module to shark_module otherwise downloads " + "the model from shark_tank.", ) p.add_argument( "--load_vmfb", default=True, action=argparse.BooleanOptionalAction, - help="attempts to load the model from a precompiled flatbuffer and compiles + saves it if not found.", + help="Attempts to load the model from a precompiled flat-buffer and compiles " + "+ saves it if not found.", ) p.add_argument( "--save_vmfb", default=False, action=argparse.BooleanOptionalAction, - help="saves the compiled flatbuffer to the local directory", + help="Saves the compiled flat-buffer to the local directory.", ) p.add_argument( "--use_tuned", default=True, action=argparse.BooleanOptionalAction, - help="Download and use the tuned version of the model if available", + help="Download and use the tuned version of the model if available.", ) p.add_argument( @@ -281,28 +291,30 @@ p.add_argument( "--scheduler", type=str, default="SharkEulerDiscrete", - help="other supported schedulers are [PNDM, DDIM, LMSDiscrete, EulerDiscrete, DPMSolverMultistep]", + help="Other supported schedulers are [DDIM, PNDM, LMSDiscrete, DPMSolverMultistep, " + "EulerDiscrete, EulerAncestralDiscrete, DEISMultistep, KDPM2AncestralDiscrete, " + "DPMSolverSinglestep, DDPM, HeunDiscrete]", ) p.add_argument( "--output_img_format", type=str, default="png", - help="specify the format in which output image is save. Supported options: jpg / png", + help="Specify the format in which output image is save. Supported options: jpg / png.", ) p.add_argument( "--output_dir", type=str, default=None, - help="Directory path to save the output images and json", + help="Directory path to save the output images and json.", ) p.add_argument( "--batch_count", type=int, default=1, - help="number of batch to be generated with random seeds in single execution", + help="Number of batch to be generated with random seeds in single execution.", ) p.add_argument( @@ -330,14 +342,14 @@ p.add_argument( "--low_cpu_mem_usage", default=False, action=argparse.BooleanOptionalAction, - help="Use the accelerate package to reduce cpu memory consumption", + help="Use the accelerate package to reduce cpu memory consumption.", ) p.add_argument( "--attention_slicing", type=str, default="none", - help="Amount of attention slicing to use (one of 'max', 'auto', 'none', or an integer)", + help="Amount of attention slicing to use (one of 'max', 'auto', 'none', or an integer).", ) p.add_argument( @@ -350,77 +362,79 @@ p.add_argument( "--use_lora", type=str, default="", - help="Use standalone LoRA weight using a HF ID or a checkpoint file (~3 MB)", + help="Use standalone LoRA weight using a HF ID or a checkpoint file (~3 MB).", ) p.add_argument( "--use_quantize", type=str, default="none", - help="""Runs the quantized version of stable diffusion model. This is currently in experimental phase. - Currently, only runs the stable-diffusion-2-1-base model in int8 quantization.""", + help="Runs the quantized version of stable diffusion model. This is currently " + "in experimental phase. Currently, only runs the stable-diffusion-2-1-base " + "model in int8 quantization.", ) p.add_argument( "--ondemand", default=False, action=argparse.BooleanOptionalAction, - help="Load and unload models for low VRAM", + help="Load and unload models for low VRAM.", ) ############################################################################## -### IREE - Vulkan supported flags +# IREE - Vulkan supported flags ############################################################################## p.add_argument( "--iree_vulkan_target_triple", type=str, default="", - help="Specify target triple for vulkan", + help="Specify target triple for vulkan.", ) p.add_argument( "--iree_metal_target_platform", type=str, default="", - help="Specify target triple for metal", + help="Specify target triple for metal.", ) p.add_argument( "--vulkan_debug_utils", default=False, action=argparse.BooleanOptionalAction, - help="Profiles vulkan device and collects the .rdc info", + help="Profiles vulkan device and collects the .rdc info.", ) p.add_argument( "--vulkan_large_heap_block_size", default="2073741824", - help="flag for setting VMA preferredLargeHeapBlockSize for vulkan device, default is 4G", + help="Flag for setting VMA preferredLargeHeapBlockSize for vulkan device, default is 4G.", ) p.add_argument( "--vulkan_validation_layers", default=False, action=argparse.BooleanOptionalAction, - help="flag for disabling vulkan validation layers when benchmarking", + help="Flag for disabling vulkan validation layers when benchmarking.", ) ############################################################################## -### Misc. Debug and Optimization flags +# Misc. Debug and Optimization flags ############################################################################## p.add_argument( "--use_compiled_scheduler", default=True, action=argparse.BooleanOptionalAction, - help="use the default scheduler precompiled into the model if available", + help="Use the default scheduler precompiled into the model if available.", ) p.add_argument( "--local_tank_cache", default="", - help="Specify where to save downloaded shark_tank artifacts. If this is not set, the default is ~/.local/shark_tank/.", + help="Specify where to save downloaded shark_tank artifacts. If this is not set, the " + "default is ~/.local/shark_tank/.", ) p.add_argument( @@ -433,133 +447,138 @@ p.add_argument( p.add_argument( "--dispatch_benchmarks", default=None, - help='dispatches to return benchamrk data on. use "All" for all, and None for none.', + help='Dispatches to return benchmark data on. Use "All" for all, and None for none.', ) p.add_argument( "--dispatch_benchmarks_dir", default="temp_dispatch_benchmarks", - help='directory where you want to store dispatch data generated with "--dispatch_benchmarks"', + help='Directory where you want to store dispatch data generated with "--dispatch_benchmarks".', ) p.add_argument( "--enable_rgp", default=False, action=argparse.BooleanOptionalAction, - help="flag for inserting debug frames between iterations for use with rgp.", + help="Flag for inserting debug frames between iterations for use with rgp.", ) p.add_argument( "--hide_steps", default=True, action=argparse.BooleanOptionalAction, - help="flag for hiding the details of iteration/sec for each step.", + help="Flag for hiding the details of iteration/sec for each step.", ) p.add_argument( "--warmup_count", type=int, default=0, - help="flag setting warmup count for clip and vae [>= 0].", + help="Flag setting warmup count for clip and vae [>= 0].", ) p.add_argument( "--clear_all", default=False, action=argparse.BooleanOptionalAction, - help="flag to clear all mlir and vmfb from common locations. Recompiling will take several minutes", + help="Flag to clear all mlir and vmfb from common locations. Recompiling will take " + "several minutes.", ) p.add_argument( "--save_metadata_to_json", default=False, action=argparse.BooleanOptionalAction, - help="flag for whether or not to save a generation information json file with the image.", + help="Flag for whether or not to save a generation information json file with the image.", ) p.add_argument( "--write_metadata_to_png", default=True, action=argparse.BooleanOptionalAction, - help="flag for whether or not to save generation information in PNG chunk text to generated images.", + help="Flag for whether or not to save generation information in PNG chunk text to " + "generated images.", ) p.add_argument( "--import_debug", default=False, action=argparse.BooleanOptionalAction, - help="if import_mlir is True, saves mlir via the debug option in shark importer. Does nothing if import_mlir is false (the default)", + help="If import_mlir is True, saves mlir via the debug option in shark importer. Does " + "nothing if import_mlir is false (the default).", ) ############################################################################## -### Web UI flags +# Web UI flags ############################################################################## p.add_argument( "--progress_bar", default=True, action=argparse.BooleanOptionalAction, - help="flag for removing the progress bar animation during image generation", + help="Flag for removing the progress bar animation during image generation.", ) p.add_argument( "--ckpt_dir", type=str, default="", - help="Path to directory where all .ckpts are stored in order to populate them in the web UI", + help="Path to directory where all .ckpts are stored in order to populate them in the web UI.", ) # TODO: replace API flag when these can be run together p.add_argument( "--ui", type=str, default="app" if os.name == "nt" else "web", - help="one of: [api, app, web]", + help="One of: [api, app, web].", ) p.add_argument( "--share", default=False, action=argparse.BooleanOptionalAction, - help="flag for generating a public URL", + help="Flag for generating a public URL.", ) p.add_argument( "--server_port", type=int, default=8080, - help="flag for setting server port", + help="Flag for setting server port.", ) p.add_argument( "--api", default=False, action=argparse.BooleanOptionalAction, - help="flag for enabling rest API", + help="Flag for enabling rest API.", ) p.add_argument( "--output_gallery", default=True, action=argparse.BooleanOptionalAction, - help="flag for removing the output gallery tab, and avoid exposing images under --output_dir in the UI", + help="Flag for removing the output gallery tab, and avoid exposing images under " + "--output_dir in the UI.", ) p.add_argument( "--output_gallery_followlinks", default=False, action=argparse.BooleanOptionalAction, - help="flag for whether the output gallery tab in the UI should follow symlinks when listing subdirectorys under --output_dir", + help="Flag for whether the output gallery tab in the UI should follow symlinks when " + "listing subdirectories under --output_dir.", ) ############################################################################## -### SD model auto-annotation flags +# SD model auto-annotation flags ############################################################################## p.add_argument( "--annotation_output", type=path_expand, default="./", - help="Directory to save the annotated mlir file", + help="Directory to save the annotated mlir file.", ) p.add_argument( @@ -573,31 +592,31 @@ p.add_argument( "--save_annotation", default=False, action=argparse.BooleanOptionalAction, - help="Save annotated mlir file", + help="Save annotated mlir file.", ) ############################################################################## -### SD model auto-tuner flags +# SD model auto-tuner flags ############################################################################## p.add_argument( "--tuned_config_dir", type=path_expand, default="./", - help="Directory to save the tuned config file", + help="Directory to save the tuned config file.", ) p.add_argument( "--num_iters", type=int, default=400, - help="Number of iterations for tuning", + help="Number of iterations for tuning.", ) p.add_argument( "--search_op", type=str, default="all", - help="Op to be optimized, options are matmul, bmm, conv and all", + help="Op to be optimized, options are matmul, bmm, conv and all.", ) diff --git a/apps/stable_diffusion/src/utils/utils.py b/apps/stable_diffusion/src/utils/utils.py index 816f024d..0cffdd2e 100644 --- a/apps/stable_diffusion/src/utils/utils.py +++ b/apps/stable_diffusion/src/utils/utils.py @@ -80,7 +80,9 @@ def _compile_module(shark_module, model_name, extra_args=[]): # Downloads the model from shark_tank and returns the shark_module. -def get_shark_model(tank_url, model_name, extra_args=[]): +def get_shark_model(tank_url, model_name, extra_args=None): + if extra_args is None: + extra_args = [] from shark.parser import shark_args # Set local shark_tank cache directory. @@ -112,13 +114,15 @@ def compile_through_fx( save_dir=tempfile.gettempdir(), debug=False, generate_vmfb=True, - extra_args=[], + extra_args=None, base_model_id=None, model_name=None, precision=None, return_mlir=False, device=None, ): + if extra_args is None: + extra_args = [] if not return_mlir and model_name is not None: vmfb_path = get_vmfb_path_name(extended_model_name) if os.path.isfile(vmfb_path): @@ -210,7 +214,8 @@ def get_device_mapping(driver, key_combination=3): 3 : (name, path) Defaults to 3. Returns: - dict: map to possible device names user can input mapped to desired combination of name/path. + dict: map to possible device names user can input mapped to desired combination + of name/path. """ from shark.iree_utils._common import iree_device_map @@ -224,7 +229,7 @@ def get_device_mapping(driver, key_combination=3): if key_combination == 2: return dev_dict["name"] if key_combination == 3: - return (dev_dict["name"], f"{driver}://{dev_dict['path']}") + return dev_dict["name"], f"{driver}://{dev_dict['path']}" # mapping driver name to default device (driver://0) device_map[f"{driver}"] = get_output_value(device_list[0]) @@ -248,7 +253,8 @@ def map_device_to_name_path(device, key_combination=3): Raises: ValueError: Returns: - str / tuple: returns the mapping str or tuple of mapping str for the device depending on key_combination value + str / tuple: returns the mapping str or tuple of mapping str for the device + depending on key_combination value """ driver = device.split("://")[0] device_map = get_device_mapping(driver, key_combination) @@ -555,7 +561,10 @@ def convert_original_vae(vae_checkpoint): for key in list(vae_checkpoint.keys()): vae_state_dict["first_stage_model." + key] = vae_checkpoint.get(key) - config_url = "https://raw.githubusercontent.com/CompVis/stable-diffusion/main/configs/stable-diffusion/v1-inference.yaml" + config_url = ( + "https://raw.githubusercontent.com/CompVis/stable-diffusion/main/configs/" + "stable-diffusion/v1-inference.yaml" + ) original_config_file = BytesIO(requests.get(config_url).content) original_config = OmegaConf.load(original_config_file) vae_config = create_vae_diffusers_config(original_config, image_size=512) @@ -676,7 +685,7 @@ def update_lora_weight(model, use_lora, model_name): # `fetch_and_update_base_model_id` is a resource utility function which -# helps maintaining mapping of the model to run with its base model. +# helps to maintain mapping of the model to run with its base model. # If `base_model` is "", then this function tries to fetch the base model # info for the `model_to_run`. def fetch_and_update_base_model_id(model_to_run, base_model=""): @@ -747,7 +756,9 @@ def get_generated_imgs_todays_subdir() -> str: # save output images and the inputs corresponding to it. -def save_output_img(output_img, img_seed, extra_info={}): +def save_output_img(output_img, img_seed, extra_info=None): + if extra_info is None: + extra_info = {} generated_imgs_path = Path( get_generated_imgs_path(), get_generated_imgs_todays_subdir() ) @@ -779,9 +790,16 @@ def save_output_img(output_img, img_seed, extra_info={}): if args.write_metadata_to_png: pngInfo.add_text( "parameters", - f"{args.prompts[0]}\nNegative prompt: {args.negative_prompts[0]}\nSteps: {args.steps}," - f"Sampler: {args.scheduler}, CFG scale: {args.guidance_scale}, Seed: {img_seed}," - f"Size: {args.width}x{args.height}, Model: {img_model}, VAE: {img_vae}, LoRA: {img_lora}", + f"{args.prompts[0]}" + f"\nNegative prompt: {args.negative_prompts[0]}" + f"\nSteps: {args.steps}," + f"Sampler: {args.scheduler}, " + f"CFG scale: {args.guidance_scale}, " + f"Seed: {img_seed}," + f"Size: {args.width}x{args.height}, " + f"Model: {img_model}, " + f"VAE: {img_vae}, " + f"LoRA: {img_lora}", ) output_img.save(out_img_path, "PNG", pnginfo=pngInfo) @@ -832,16 +850,27 @@ def save_output_img(output_img, img_seed, extra_info={}): def get_generation_text_info(seeds, device): text_output = f"prompt={args.prompts}" text_output += f"\nnegative prompt={args.negative_prompts}" - text_output += f"\nmodel_id={args.hf_model_id}, ckpt_loc={args.ckpt_loc}" - text_output += f"\nscheduler={args.scheduler}, device={device}" - text_output += f"\nsteps={args.steps}, guidance_scale={args.guidance_scale}, seed={seeds}" - text_output += f"\nsize={args.height}x{args.width}, batch_count={args.batch_count}, batch_size={args.batch_size}, max_length={args.max_length}" + text_output += ( + f"\nmodel_id={args.hf_model_id}, " f"ckpt_loc={args.ckpt_loc}" + ) + text_output += f"\nscheduler={args.scheduler}, " f"device={device}" + text_output += ( + f"\nsteps={args.steps}, " + f"guidance_scale={args.guidance_scale}, " + f"seed={seeds}" + ) + text_output += ( + f"\nsize={args.height}x{args.width}, " + f"batch_count={args.batch_count}, " + f"batch_size={args.batch_size}, " + f"max_length={args.max_length}" + ) return text_output -# For stencil, the input image can be of any size but we need to ensure that -# it conforms with our model contraints :- +# For stencil, the input image can be of any size, but we need to ensure that +# it conforms with our model constraints :- # Both width and height should be in the range of [128, 768] and multiple of 8. # This utility function performs the transformation on the input image while # also maintaining the aspect ratio before sending it to the stencil pipeline. diff --git a/apps/stable_diffusion/web/index.py b/apps/stable_diffusion/web/index.py index 7cfcf3cc..d7ec0202 100644 --- a/apps/stable_diffusion/web/index.py +++ b/apps/stable_diffusion/web/index.py @@ -26,9 +26,10 @@ def launch_app(address): window = Tk() - # getting screen width and height of display - width = window.winfo_screenwidth() - height = window.winfo_screenheight() + # get screen width and height of display and make it more reasonably + # sized as we aren't making it full-screen or maximized + width = int(window.winfo_screenwidth() * 0.81) + height = int(window.winfo_screenheight() * 0.91) webview.create_window( "SHARK AI Studio", url=address, diff --git a/apps/stable_diffusion/web/ui/img2img_ui.py b/apps/stable_diffusion/web/ui/img2img_ui.py index ea960364..a18599a8 100644 --- a/apps/stable_diffusion/web/ui/img2img_ui.py +++ b/apps/stable_diffusion/web/ui/img2img_ui.py @@ -104,7 +104,8 @@ def img2img_inf( if not hf_model_id: return ( None, - "Please provide either custom model or huggingface model ID, both must not be empty", + "Please provide either custom model or huggingface model ID, " + "both must not be empty.", ) if "civitai" in hf_model_id: args.ckpt_loc = hf_model_id @@ -308,7 +309,9 @@ def img2img_api( InputData: dict, ): print( - f'Prompt: {InputData["prompt"]}, Negative Prompt: {InputData["negative_prompt"]}, Seed: {InputData["seed"]}' + f'Prompt: {InputData["prompt"]}, ' + f'Negative Prompt: {InputData["negative_prompt"]}, ' + f'Seed: {InputData["seed"]}.' ) init_image = decode_base64_to_image(InputData["init_images"][0]) res = img2img_inf( @@ -367,8 +370,14 @@ with gr.Blocks(title="Image-to-Image") as img2img_web: with gr.Row(): with gr.Column(scale=1, min_width=600): with gr.Row(): + # janky fix for overflowing text + i2i_model_info = (str(get_custom_model_path())).replace( + "\\", "\n\\" + ) + i2i_model_info = f"Custom Model Path: {i2i_model_info}" img2img_custom_model = gr.Dropdown( - label=f"Models (Custom Model path: {get_custom_model_path()})", + label=f"Models", + info=i2i_model_info, elem_id="custom_model", value=os.path.basename(args.ckpt_loc) if args.ckpt_loc @@ -379,13 +388,22 @@ with gr.Blocks(title="Image-to-Image") as img2img_web: ) img2img_hf_model_id = gr.Textbox( elem_id="hf_model_id", - placeholder="Select 'None' in the Models dropdown on the left and enter model ID here e.g: SG161222/Realistic_Vision_V1.3, https://civitai.com/api/download/models/15236", + placeholder="Select 'None' in the Models dropdown " + "on the left and enter model ID here " + "e.g: SG161222/Realistic_Vision_V1.3, " + "https://civitai.com/api/download/models/15236", value="", label="HuggingFace Model ID or Civitai model download URL", lines=3, ) + # janky fix for overflowing text + i2i_vae_info = (str(get_custom_model_path("vae"))).replace( + "\\", "\n\\" + ) + i2i_vae_info = f"VAE Path: {i2i_vae_info}" custom_vae = gr.Dropdown( - label=f"Custom Vae Models (Path: {get_custom_model_path('vae')})", + label=f"Custom VAE Models", + info=i2i_vae_info, elem_id="custom_model", value=os.path.basename(args.custom_vae) if args.custom_vae @@ -397,13 +415,13 @@ with gr.Blocks(title="Image-to-Image") as img2img_web: prompt = gr.Textbox( label="Prompt", value=args.prompts[0], - lines=1, + lines=2, elem_id="prompt_box", ) negative_prompt = gr.Textbox( label="Negative Prompt", value=args.negative_prompts[0], - lines=1, + lines=2, elem_id="negative_prompt_box", ) @@ -475,15 +493,24 @@ with gr.Blocks(title="Image-to-Image") as img2img_web: with gr.Accordion(label="LoRA Options", open=False): with gr.Row(): + # janky fix for overflowing text + i2i_lora_info = ( + str(get_custom_model_path("lora")) + ).replace("\\", "\n\\") + i2i_lora_info = f"LoRA Path: {i2i_lora_info}" lora_weights = gr.Dropdown( - label=f"Standlone LoRA weights (Path: {get_custom_model_path('lora')})", + label=f"Standalone LoRA Weights", + info=i2i_lora_info, elem_id="lora_weights", value="None", choices=["None"] + get_custom_model_files("lora"), ) lora_hf_id = gr.Textbox( elem_id="lora_hf_id", - placeholder="Select 'None' in the Standlone LoRA weights dropdown on the left if you want to use a standalone HuggingFace model ID for LoRA here e.g: sayakpaul/sd-model-finetuned-lora-t4", + placeholder="Select 'None' in the Standalone LoRA weights dropdown " + "on the left if you want to use a standalone " + "HuggingFace model ID for LoRA here " + "e.g: sayakpaul/sd-model-finetuned-lora-t4", value="", label="HuggingFace Model ID", lines=3, diff --git a/apps/stable_diffusion/web/ui/inpaint_ui.py b/apps/stable_diffusion/web/ui/inpaint_ui.py index 6bd5a3c3..37ec0e70 100644 --- a/apps/stable_diffusion/web/ui/inpaint_ui.py +++ b/apps/stable_diffusion/web/ui/inpaint_ui.py @@ -92,7 +92,8 @@ def inpaint_inf( if not hf_model_id: return ( None, - "Please provide either custom model or huggingface model ID, both must not be empty", + "Please provide either custom model or huggingface model ID, " + "both must not be empty.", ) if "civitai" in hf_model_id: args.ckpt_loc = hf_model_id @@ -258,7 +259,9 @@ def inpaint_api( InputData: dict, ): print( - f'Prompt: {InputData["prompt"]}, Negative Prompt: {InputData["negative_prompt"]}, Seed: {InputData["seed"]}' + f'Prompt: {InputData["prompt"]}, ' + f'Negative Prompt: {InputData["negative_prompt"]}, ' + f'Seed: {InputData["seed"]}.' ) init_image = decode_base64_to_image(InputData["image"]) mask = decode_base64_to_image(InputData["mask"]) @@ -316,8 +319,16 @@ with gr.Blocks(title="Inpainting") as inpaint_web: with gr.Row(): with gr.Column(scale=1, min_width=600): with gr.Row(): + # janky fix for overflowing text + inpaint_model_info = ( + str(get_custom_model_path()) + ).replace("\\", "\n\\") + inpaint_model_info = ( + f"Custom Model Path: {inpaint_model_info}" + ) inpaint_custom_model = gr.Dropdown( - label=f"Models (Custom Model path: {get_custom_model_path()})", + label=f"Models", + info=inpaint_model_info, elem_id="custom_model", value=os.path.basename(args.ckpt_loc) if args.ckpt_loc @@ -330,13 +341,22 @@ with gr.Blocks(title="Inpainting") as inpaint_web: ) inpaint_hf_model_id = gr.Textbox( elem_id="hf_model_id", - placeholder="Select 'None' in the Models dropdown on the left and enter model ID here e.g: ghunkins/stable-diffusion-liberty-inpainting, https://civitai.com/api/download/models/3433", + placeholder="Select 'None' in the Models dropdown " + "on the left and enter model ID here " + "e.g: ghunkins/stable-diffusion-liberty-inpainting, " + "https://civitai.com/api/download/models/3433", value="", label="HuggingFace Model ID or Civitai model download URL", lines=3, ) + # janky fix for overflowing text + inpaint_vae_info = ( + str(get_custom_model_path("vae")) + ).replace("\\", "\n\\") + inpaint_vae_info = f"VAE Path: {inpaint_vae_info}" custom_vae = gr.Dropdown( - label=f"Custom Vae Models (Path: {get_custom_model_path('vae')})", + label=f"Custom VAE Models", + info=inpaint_vae_info, elem_id="custom_model", value=os.path.basename(args.custom_vae) if args.custom_vae @@ -348,13 +368,13 @@ with gr.Blocks(title="Inpainting") as inpaint_web: prompt = gr.Textbox( label="Prompt", value=args.prompts[0], - lines=1, + lines=2, elem_id="prompt_box", ) negative_prompt = gr.Textbox( label="Negative Prompt", value=args.negative_prompts[0], - lines=1, + lines=2, elem_id="negative_prompt_box", ) @@ -367,15 +387,24 @@ with gr.Blocks(title="Inpainting") as inpaint_web: with gr.Accordion(label="LoRA Options", open=False): with gr.Row(): + # janky fix for overflowing text + inpaint_lora_info = ( + str(get_custom_model_path("lora")) + ).replace("\\", "\n\\") + inpaint_lora_info = f"LoRA Path: {inpaint_lora_info}" lora_weights = gr.Dropdown( - label=f"Standlone LoRA weights (Path: {get_custom_model_path('lora')})", + label=f"Standalone LoRA Weights", + info=inpaint_lora_info, elem_id="lora_weights", value="None", choices=["None"] + get_custom_model_files("lora"), ) lora_hf_id = gr.Textbox( elem_id="lora_hf_id", - placeholder="Select 'None' in the Standlone LoRA weights dropdown on the left if you want to use a standalone HuggingFace model ID for LoRA here e.g: sayakpaul/sd-model-finetuned-lora-t4", + placeholder="Select 'None' in the Standalone LoRA weights dropdown " + "on the left if you want to use a standalone " + "HuggingFace model ID for LoRA here " + "e.g: sayakpaul/sd-model-finetuned-lora-t4", value="", label="HuggingFace Model ID", lines=3, diff --git a/apps/stable_diffusion/web/ui/lora_train_ui.py b/apps/stable_diffusion/web/ui/lora_train_ui.py index 9109c111..0439a81c 100644 --- a/apps/stable_diffusion/web/ui/lora_train_ui.py +++ b/apps/stable_diffusion/web/ui/lora_train_ui.py @@ -31,8 +31,16 @@ with gr.Blocks(title="Lora Training") as lora_train_web: with gr.Row(): with gr.Column(scale=10): with gr.Row(): + # janky fix for overflowing text + train_lora_model_info = ( + str(get_custom_model_path()) + ).replace("\\", "\n\\") + train_lora_model_info = ( + f"Custom Model Path: {train_lora_model_info}" + ) custom_model = gr.Dropdown( - label=f"Models (Custom Model path: {get_custom_model_path()})", + label=f"Models", + info=train_lora_model_info, elem_id="custom_model", value=os.path.basename(args.ckpt_loc) if args.ckpt_loc @@ -43,22 +51,33 @@ with gr.Blocks(title="Lora Training") as lora_train_web: ) hf_model_id = gr.Textbox( elem_id="hf_model_id", - placeholder="Select 'None' in the Models dropdown on the left and enter model ID here e.g: SG161222/Realistic_Vision_V1.3", + placeholder="Select 'None' in the Models dropdown " + "on the left and enter model ID here " + "e.g: SG161222/Realistic_Vision_V1.3", value="", label="HuggingFace Model ID", lines=3, ) with gr.Row(): + # janky fix for overflowing text + train_lora_info = ( + str(get_custom_model_path("lora")) + ).replace("\\", "\n\\") + train_lora_info = f"LoRA Path: {train_lora_info}" lora_weights = gr.Dropdown( - label=f"Standlone LoRA weights to initialize weights (Path: {get_custom_model_path('lora')})", + label=f"Standalone LoRA weights to initialize weights", + info=train_lora_info, elem_id="lora_weights", value="None", choices=["None"] + get_custom_model_files("lora"), ) lora_hf_id = gr.Textbox( elem_id="lora_hf_id", - placeholder="Select 'None' in the Standlone LoRA weights dropdown on the left if you want to use a standalone HuggingFace model ID for LoRA here e.g: sayakpaul/sd-model-finetuned-lora-t4", + placeholder="Select 'None' in the Standalone LoRA weights " + "dropdown on the left if you want to use a " + "standalone HuggingFace model ID for LoRA here " + "e.g: sayakpaul/sd-model-finetuned-lora-t4", value="", label="HuggingFace Model ID to initialize weights", lines=3, @@ -74,7 +93,7 @@ with gr.Blocks(title="Lora Training") as lora_train_web: prompt = gr.Textbox( label="Prompt", value=args.prompts[0], - lines=1, + lines=2, elem_id="prompt_box", ) with gr.Accordion(label="Advanced Options", open=False): diff --git a/apps/stable_diffusion/web/ui/outpaint_ui.py b/apps/stable_diffusion/web/ui/outpaint_ui.py index 9439dd15..3938c74b 100644 --- a/apps/stable_diffusion/web/ui/outpaint_ui.py +++ b/apps/stable_diffusion/web/ui/outpaint_ui.py @@ -29,7 +29,6 @@ from apps.stable_diffusion.src.utils import ( ) from apps.stable_diffusion.web.utils.common_label_calc import status_label - # set initial values of iree_vulkan_target_triple, use_tuned and import_mlir. init_iree_vulkan_target_triple = args.iree_vulkan_target_triple init_use_tuned = args.use_tuned @@ -92,7 +91,8 @@ def outpaint_inf( if not hf_model_id: return ( None, - "Please provide either custom model or huggingface model ID, both must not be empty", + "Please provide either custom model or huggingface model ID, " + "both must not be empty.", ) if "civitai" in hf_model_id: args.ckpt_loc = hf_model_id @@ -325,8 +325,16 @@ with gr.Blocks(title="Outpainting") as outpaint_web: with gr.Row(): with gr.Column(scale=1, min_width=600): with gr.Row(): + # janky fix for overflowing text + outpaint_model_info = ( + str(get_custom_model_path()) + ).replace("\\", "\n\\") + outpaint_model_info = ( + f"Custom Model Path: {outpaint_model_info}" + ) outpaint_custom_model = gr.Dropdown( - label=f"Models (Custom Model path: {get_custom_model_path()})", + label=f"Models", + info=outpaint_model_info, elem_id="custom_model", value=os.path.basename(args.ckpt_loc) if args.ckpt_loc @@ -339,13 +347,22 @@ with gr.Blocks(title="Outpainting") as outpaint_web: ) outpaint_hf_model_id = gr.Textbox( elem_id="hf_model_id", - placeholder="Select 'None' in the Models dropdown on the left and enter model ID here e.g: ghunkins/stable-diffusion-liberty-inpainting, https://civitai.com/api/download/models/3433", + placeholder="Select 'None' in the Models dropdown " + "on the left and enter model ID here " + "e.g: ghunkins/stable-diffusion-liberty-inpainting, " + "https://civitai.com/api/download/models/3433", value="", label="HuggingFace Model ID or Civitai model download URL", lines=3, ) + # janky fix for overflowing text + outpaint_vae_info = ( + str(get_custom_model_path("vae")) + ).replace("\\", "\n\\") + outpaint_vae_info = f"VAE Path: {outpaint_vae_info}" custom_vae = gr.Dropdown( - label=f"Custom Vae Models (Path: {get_custom_model_path('vae')})", + label=f"Custom VAE Models", + info=outpaint_vae_info, elem_id="custom_model", value=os.path.basename(args.custom_vae) if args.custom_vae @@ -357,13 +374,13 @@ with gr.Blocks(title="Outpainting") as outpaint_web: prompt = gr.Textbox( label="Prompt", value=args.prompts[0], - lines=1, + lines=2, elem_id="prompt_box", ) negative_prompt = gr.Textbox( label="Negative Prompt", value=args.negative_prompts[0], - lines=1, + lines=2, elem_id="negative_prompt_box", ) @@ -373,15 +390,24 @@ with gr.Blocks(title="Outpainting") as outpaint_web: with gr.Accordion(label="LoRA Options", open=False): with gr.Row(): + # janky fix for overflowing text + outpaint_lora_info = ( + str(get_custom_model_path("lora")) + ).replace("\\", "\n\\") + outpaint_lora_info = f"LoRA Path: {outpaint_lora_info}" lora_weights = gr.Dropdown( - label=f"Standlone LoRA weights (Path: {get_custom_model_path('lora')})", + label=f"Standalone LoRA Weights", + info=outpaint_lora_info, elem_id="lora_weights", value="None", choices=["None"] + get_custom_model_files("lora"), ) lora_hf_id = gr.Textbox( elem_id="lora_hf_id", - placeholder="Select 'None' in the Standlone LoRA weights dropdown on the left if you want to use a standalone HuggingFace model ID for LoRA here e.g: sayakpaul/sd-model-finetuned-lora-t4", + placeholder="Select 'None' in the Standalone LoRA weights dropdown " + "on the left if you want to use a " + "standalone HuggingFace model ID for LoRA here " + "e.g: sayakpaul/sd-model-finetuned-lora-t4", value="", label="HuggingFace Model ID", lines=3, diff --git a/apps/stable_diffusion/web/ui/outputgallery_ui.py b/apps/stable_diffusion/web/ui/outputgallery_ui.py index 687afd87..42c0fd75 100644 --- a/apps/stable_diffusion/web/ui/outputgallery_ui.py +++ b/apps/stable_diffusion/web/ui/outputgallery_ui.py @@ -192,7 +192,7 @@ with gr.Blocks() as outputgallery_web: ] def on_refresh(current_subdir: str) -> list: - # get an up to date subdirectory list + # get an up-to-date subdirectory list refreshed_subdirs = output_subdirs() # get the images using either the current subdirectory or the most recent valid one new_subdir = ( diff --git a/apps/stable_diffusion/web/ui/stablelm_ui.py b/apps/stable_diffusion/web/ui/stablelm_ui.py index 07e9d7b6..4d77f5a8 100644 --- a/apps/stable_diffusion/web/ui/stablelm_ui.py +++ b/apps/stable_diffusion/web/ui/stablelm_ui.py @@ -7,12 +7,15 @@ from transformers import ( ) from apps.stable_diffusion.web.ui.utils import available_devices -start_message = """<|SYSTEM|># StableLM Tuned (Alpha version) -- StableLM is a helpful and harmless open-source AI language model developed by StabilityAI. -- StableLM is excited to be able to help the user, but will refuse to do anything that could be considered harmful to the user. -- StableLM is more than just an information source, StableLM is also able to write poetry, short stories, and make jokes. -- StableLM will refuse to participate in anything that could harm a human. -""" +start_message = ( + "<|SYSTEM|># StableLM Tuned (Alpha version)- StableLM is a helpful and " + "harmless open-source AI language model developed by StabilityAI." + "\n- StableLM is excited to be able to help the user, but will refuse to do " + "anything that could be considered harmful to the user." + "\n- StableLM is more than just an information source, StableLM is also able to " + "write poetry, short stories, and make jokes." + "\n- StableLM will refuse to participate in anything that could harm a human." +) def user(message, history): @@ -25,7 +28,11 @@ sharded_model = 0 vicuna_model = 0 -start_message_vicuna = "A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions.\n" +start_message_vicuna = ( + "A chat between a curious user and an artificial intelligence assistant. " + "The assistant gives helpful, detailed, and polite answers to the user's " + "questions.\n" +) past_key_values = None @@ -104,7 +111,7 @@ def chat(curr_system_message, history, model, device, precision): # print(new_text) partial_text += new_text history[-1][1] = partial_text - # Yield an empty string to cleanup the message textbox and the updated conversation history + # Yield an empty string to clean up the message textbox and the updated conversation history yield history return words_list diff --git a/apps/stable_diffusion/web/ui/txt2img_ui.py b/apps/stable_diffusion/web/ui/txt2img_ui.py index 03185f5b..f85324ea 100644 --- a/apps/stable_diffusion/web/ui/txt2img_ui.py +++ b/apps/stable_diffusion/web/ui/txt2img_ui.py @@ -87,7 +87,8 @@ def txt2img_inf( if not hf_model_id: return ( None, - "Please provide either custom model or huggingface model ID, both must not be empty", + "Please provide either custom model or huggingface model ID, both " + "must not be empty", ) if "civitai" in hf_model_id: args.ckpt_loc = hf_model_id @@ -238,7 +239,9 @@ def txt2img_api( InputData: dict, ): print( - f'Prompt: {InputData["prompt"]}, Negative Prompt: {InputData["negative_prompt"]}, Seed: {InputData["seed"]}' + f'Prompt: {InputData["prompt"]}, ' + f'Negative Prompt: {InputData["negative_prompt"]}, ' + f'Seed: {InputData["seed"]}.' ) res = txt2img_inf( InputData["prompt"], @@ -293,8 +296,16 @@ with gr.Blocks(title="Text-to-Image") as txt2img_web: with gr.Row(): with gr.Column(scale=10): with gr.Row(): + # janky fix for overflowing text + t2i_model_info = ( + str(get_custom_model_path()) + ).replace("\\", "\n\\") + t2i_model_info = ( + f"Custom Model Path: {t2i_model_info}" + ) txt2img_custom_model = gr.Dropdown( - label=f"Models (Custom Model path: {get_custom_model_path()})", + label=f"Models", + info=t2i_model_info, elem_id="custom_model", value=os.path.basename(args.ckpt_loc) if args.ckpt_loc @@ -305,13 +316,20 @@ with gr.Blocks(title="Text-to-Image") as txt2img_web: ) txt2img_hf_model_id = gr.Textbox( elem_id="hf_model_id", - placeholder="Select 'None' in the dropdown on the left and enter model ID here", + placeholder="Select 'None' in the dropdown on the left and " + "enter model ID here.", value="", - label="HuggingFace Model ID or Civitai model download URL", + label="HuggingFace Model ID or Civitai model download URL.", lines=3, ) + # janky fix for overflowing text + t2i_vae_info = ( + str(get_custom_model_path("vae")) + ).replace("\\", "\n\\") + t2i_vae_info = f"VAE Path: {t2i_vae_info}" custom_vae = gr.Dropdown( - label=f"Custom Vae Models (Path: {get_custom_model_path('vae')})", + label=f"VAE Models", + info=t2i_vae_info, elem_id="custom_model", value=os.path.basename(args.custom_vae) if args.custom_vae @@ -332,26 +350,35 @@ with gr.Blocks(title="Text-to-Image") as txt2img_web: prompt = gr.Textbox( label="Prompt", value=args.prompts[0], - lines=1, + lines=2, elem_id="prompt_box", ) negative_prompt = gr.Textbox( label="Negative Prompt", value=args.negative_prompts[0], - lines=1, + lines=2, elem_id="negative_prompt_box", ) with gr.Accordion(label="LoRA Options", open=False): with gr.Row(): + # janky fix for overflowing text + t2i_lora_info = ( + str(get_custom_model_path("lora")) + ).replace("\\", "\n\\") + t2i_lora_info = f"LoRA Path: {t2i_lora_info}" lora_weights = gr.Dropdown( - label=f"Standlone LoRA weights (Path: {get_custom_model_path('lora')})", + label=f"Standalone LoRA Weights", + info=t2i_lora_info, elem_id="lora_weights", value="None", choices=["None"] + get_custom_model_files("lora"), ) lora_hf_id = gr.Textbox( elem_id="lora_hf_id", - placeholder="Select 'None' in the Standlone LoRA weights dropdown on the left if you want to use a standalone HuggingFace model ID for LoRA here e.g: sayakpaul/sd-model-finetuned-lora-t4", + placeholder="Select 'None' in the Standalone LoRA weights dropdown " + "on the left if you want to use a standalone " + "HuggingFace model ID for LoRA here e.g: " + "sayakpaul/sd-model-finetuned-lora-t4", value="", label="HuggingFace Model ID", lines=3, diff --git a/apps/stable_diffusion/web/ui/upscaler_ui.py b/apps/stable_diffusion/web/ui/upscaler_ui.py index f36138f6..419d78cf 100644 --- a/apps/stable_diffusion/web/ui/upscaler_ui.py +++ b/apps/stable_diffusion/web/ui/upscaler_ui.py @@ -88,7 +88,8 @@ def upscaler_inf( if not hf_model_id: return ( None, - "Please provide either custom model or huggingface model ID, both must not be empty", + "Please provide either custom model or huggingface model ID, both must not be " + "empty.", ) if "civitai" in hf_model_id: args.ckpt_loc = hf_model_id @@ -325,8 +326,16 @@ with gr.Blocks(title="Upscaler") as upscaler_web: with gr.Row(): with gr.Column(scale=1, min_width=600): with gr.Row(): + # janky fix for overflowing text + upscaler_model_info = ( + str(get_custom_model_path()) + ).replace("\\", "\n\\") + upscaler_model_info = ( + f"Custom Model Path: {upscaler_model_info}" + ) upscaler_custom_model = gr.Dropdown( - label=f"Models (Custom Model path: {get_custom_model_path()})", + label=f"Models", + info=upscaler_model_info, elem_id="custom_model", value=os.path.basename(args.ckpt_loc) if args.ckpt_loc @@ -339,13 +348,22 @@ with gr.Blocks(title="Upscaler") as upscaler_web: ) upscaler_hf_model_id = gr.Textbox( elem_id="hf_model_id", - placeholder="Select 'None' in the Models dropdown on the left and enter model ID here e.g: SG161222/Realistic_Vision_V1.3, https://civitai.com/api/download/models/15236", + placeholder="Select 'None' in the Models dropdown " + "on the left and enter model ID here " + "e.g: SG161222/Realistic_Vision_V1.3, " + "https://civitai.com/api/download/models/15236", value="", label="HuggingFace Model ID or Civitai model download URL", lines=3, ) + # janky fix for overflowing text + upscaler_vae_info = ( + str(get_custom_model_path("vae")) + ).replace("\\", "\n\\") + upscaler_vae_info = f"VAE Path: {upscaler_vae_info}" custom_vae = gr.Dropdown( - label=f"Custom Vae Models (Path: {get_custom_model_path('vae')})", + label=f"Custom VAE Models", + info=upscaler_vae_info, elem_id="custom_model", value=os.path.basename(args.custom_vae) if args.custom_vae @@ -357,13 +375,13 @@ with gr.Blocks(title="Upscaler") as upscaler_web: prompt = gr.Textbox( label="Prompt", value=args.prompts[0], - lines=1, + lines=2, elem_id="prompt_box", ) negative_prompt = gr.Textbox( label="Negative Prompt", value=args.negative_prompts[0], - lines=1, + lines=2, elem_id="negative_prompt_box", ) @@ -373,15 +391,24 @@ with gr.Blocks(title="Upscaler") as upscaler_web: with gr.Accordion(label="LoRA Options", open=False): with gr.Row(): + # janky fix for overflowing text + upscaler_lora_info = ( + str(get_custom_model_path("lora")) + ).replace("\\", "\n\\") + upscaler_lora_info = f"LoRA Path: {upscaler_lora_info}" lora_weights = gr.Dropdown( - label=f"Standlone LoRA weights (Path: {get_custom_model_path('lora')})", + label=f"Standalone LoRA weights (Path: {get_custom_model_path('lora')})", + info=upscaler_lora_info, elem_id="lora_weights", value="None", choices=["None"] + get_custom_model_files("lora"), ) lora_hf_id = gr.Textbox( elem_id="lora_hf_id", - placeholder="Select 'None' in the Standlone LoRA weights dropdown on the left if you want to use a standalone HuggingFace model ID for LoRA here e.g: sayakpaul/sd-model-finetuned-lora-t4", + placeholder="Select 'None' in the Standalone LoRA weights dropdown " + "on the left if you want to use a standalone " + "HuggingFace model ID for LoRA here " + "e.g: sayakpaul/sd-model-finetuned-lora-t4", value="", label="HuggingFace Model ID", lines=3, diff --git a/apps/stable_diffusion/web/ui/utils.py b/apps/stable_diffusion/web/ui/utils.py index 0e9189cc..ba16e786 100644 --- a/apps/stable_diffusion/web/ui/utils.py +++ b/apps/stable_diffusion/web/ui/utils.py @@ -41,6 +41,11 @@ scheduler_list_cpu_only = [ "DPMSolverMultistep", "EulerDiscrete", "EulerAncestralDiscrete", + "DEISMultistep", + "KDPM2AncestralDiscrete", + "DPMSolverSinglestep", + "DDPM", + "HeunDiscrete", ] scheduler_list = scheduler_list_cpu_only + [ "SharkEulerDiscrete", @@ -50,6 +55,7 @@ predefined_models = [ "Linaqruf/anything-v3.0", "prompthero/openjourney", "wavymulder/Analog-Diffusion", + "xzuyn/PhotoMerge", "stabilityai/stable-diffusion-2-1", "stabilityai/stable-diffusion-2-1-base", "CompVis/stable-diffusion-v1-4", @@ -58,6 +64,7 @@ predefined_models = [ predefined_paint_models = [ "runwayml/stable-diffusion-inpainting", "stabilityai/stable-diffusion-2-inpainting", + "xzuyn/PhotoMerge-inpainting", ] predefined_upscaler_models = [ "stabilityai/stable-diffusion-x4-upscaler",