diff --git a/apps/stable_diffusion/src/utils/utils.py b/apps/stable_diffusion/src/utils/utils.py index 801a33e17a..1fcc03db09 100644 --- a/apps/stable_diffusion/src/utils/utils.py +++ b/apps/stable_diffusion/src/utils/utils.py @@ -777,7 +777,7 @@ def save_output_img(output_img, img_seed, extra_info={}): "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"Size: {args.width}x{args.height}, Model: {img_model}, VAE: {img_vae}, LoRA: {img_lora}", ) output_img.save(out_img_path, "PNG", pnginfo=pngInfo) diff --git a/apps/stable_diffusion/web/ui/txt2img_ui.py b/apps/stable_diffusion/web/ui/txt2img_ui.py index db5f42efdb..00585046cf 100644 --- a/apps/stable_diffusion/web/ui/txt2img_ui.py +++ b/apps/stable_diffusion/web/ui/txt2img_ui.py @@ -555,7 +555,7 @@ def txt2img_api( txt2img_hf_model_id, lora_weights, lora_hf_id, - custom_vae + custom_vae, ], outputs=[ txt2img_png_info_img, @@ -571,6 +571,6 @@ def txt2img_api( txt2img_hf_model_id, lora_weights, lora_hf_id, - custom_vae + custom_vae, ], ) diff --git a/apps/stable_diffusion/web/utils/metadata/png_metadata.py b/apps/stable_diffusion/web/utils/metadata/png_metadata.py index 8c22952699..51a92f07d6 100644 --- a/apps/stable_diffusion/web/utils/metadata/png_metadata.py +++ b/apps/stable_diffusion/web/utils/metadata/png_metadata.py @@ -61,9 +61,9 @@ def parse_generation_parameters(x: str): return res + def try_find_model_base_from_png_metadata( - file: str, - folder: str = "models" + file: str, folder: str = "models" ) -> str: custom = "" @@ -73,14 +73,16 @@ def try_find_model_base_from_png_metadata( # Check for the file name match with one of the local ckpt or safetensors files if Path(get_custom_model_pathfile(file + ".ckpt", folder)).is_file(): custom = file + ".ckpt" - if Path(get_custom_model_pathfile(file + ".safetensors", folder)).is_file(): + if Path( + get_custom_model_pathfile(file + ".safetensors", folder) + ).is_file(): custom = file + ".safetensors" return custom + def find_model_from_png_metadata( - key: str, - metadata: dict[str, str | int] + key: str, metadata: dict[str, str | int] ) -> tuple[str, str]: png_hf_id = "" png_custom = "" @@ -105,8 +107,7 @@ def find_model_from_png_metadata( def find_vae_from_png_metadata( - key: str, - metadata: dict[str, str | int] + key: str, metadata: dict[str, str | int] ) -> str: vae_custom = "" @@ -118,9 +119,9 @@ def find_vae_from_png_metadata( return vae_custom + def find_lora_from_png_metadata( - key: str, - metadata: dict[str, str | int] + key: str, metadata: dict[str, str | int] ) -> tuple[str, str]: lora_hf_id = "" lora_custom = "" @@ -136,6 +137,7 @@ def find_lora_from_png_metadata( return lora_custom, lora_hf_id + def import_png_metadata( pil_data, prompt, @@ -150,14 +152,18 @@ def import_png_metadata( hf_model_id, custom_lora, hf_lora_id, - custom_vae + custom_vae, ): try: png_info = pil_data.info["parameters"] metadata = parse_generation_parameters(png_info) - (png_custom_model, png_hf_model_id) = find_model_from_png_metadata("Model", metadata) - (lora_custom_model, lora_hf_model_id) = find_lora_from_png_metadata("LoRA", metadata) + (png_custom_model, png_hf_model_id) = find_model_from_png_metadata( + "Model", metadata + ) + (lora_custom_model, lora_hf_model_id) = find_lora_from_png_metadata( + "LoRA", metadata + ) vae_custom_model = find_vae_from_png_metadata("VAE", metadata) negative_prompt = metadata["Negative prompt"] @@ -214,5 +220,5 @@ def import_png_metadata( hf_model_id, custom_lora, hf_lora_id, - custom_vae + custom_vae, )