From 5b338548c4196811cdcae4feb9e947c8beae652b Mon Sep 17 00:00:00 2001 From: tin2tin Date: Sat, 26 Aug 2023 15:33:19 +0200 Subject: [PATCH] Add vae fix & inpaint for images --- __init__.py | 186 +++++++++++++++++++++++++++++++++++++++++----------- 1 file changed, 147 insertions(+), 39 deletions(-) diff --git a/__init__.py b/__init__.py index 099efca..8215b22 100644 --- a/__init__.py +++ b/__init__.py @@ -326,7 +326,7 @@ def process_video(input_video_path, output_video_path): Image.MAX_IMAGE_PIXELS = None import cv2 import shutil - + scene = bpy.context.scene movie_x = scene.generate_movie_x @@ -926,6 +926,10 @@ class SEQEUNCER_PT_generate_ai(Panel): # UI bl_region_type = "UI" bl_category = "Generative AI" + @classmethod + def poll(cls, context): + return context.area.type == 'SEQUENCE_EDITOR' + def draw(self, context): preferences = context.preferences addon_prefs = preferences.addons[__name__].preferences @@ -938,24 +942,38 @@ class SEQEUNCER_PT_generate_ai(Panel): # UI input = scene.input_strips layout = self.layout - col = layout.column(align=True) + col = layout.column(align=False) col.use_property_split = True col.use_property_decorate = False - col.scale_y = 1.2 - col.prop(context.scene, "generate_movie_prompt", text="", icon="ADD") + col = col.box() + col = col.column() + + col.prop(context.scene, "input_strips", text="Input") + + if input == "input_strips": + col.prop(context.scene, "image_power", text="Strip Power") + if type == "image": + col.prop_search(scene, "inpaint_selected_strip", scene.sequence_editor, "sequences", text="Inpaint Mask", icon='SEQ_STRIP_DUPLICATE') + + #layout = self.layout + col = col.column(align=True) + col.use_property_split = True + col.use_property_decorate = False + col.prop(context.scene, "generate_movie_prompt", text="Prompt", icon="ADD") if type == "audio" and audio_model_card == "bark": pass else: col.prop( - context.scene, "generate_movie_negative_prompt", text="", icon="REMOVE" - ) + context.scene, "generate_movie_negative_prompt", text="Negative Prompt", icon="REMOVE") - col.prop(context.scene, "generatorai_styles", text="Style") - - layout = self.layout + layout = col.column() layout.use_property_split = True layout.use_property_decorate = False + col = layout.column(align=True) + + if type != "audio": + col.prop(context.scene, "generatorai_styles", text="Style") if type == "movie" or type == "image": col = layout.column(align=True) @@ -978,45 +996,51 @@ class SEQEUNCER_PT_generate_ai(Panel): # UI col.prop(context.scene, "movie_num_inference_steps", text="Quality Steps") col.prop(context.scene, "movie_num_guidance", text="Word Power") - col = layout.column() + col = col.column() row = col.row(align=True) sub_row = row.row(align=True) sub_row.prop(context.scene, "movie_num_seed", text="Seed") row.prop(context.scene, "movie_use_random", text="", icon="QUESTION") sub_row.active = not context.scene.movie_use_random + layout = self.layout + layout.use_property_split = True + layout.use_property_decorate = False + col = layout.column(align=True) + col = col.box() + + col.prop(context.scene, "generatorai_typeselect", text="Output") + if type == "movie" and ( movie_model_card == "cerspense/zeroscope_v2_dark_30x448x256" or movie_model_card == "cerspense/zeroscope_v2_576w" or movie_model_card == "cerspense/zeroscope_v2_XL" ): - col = layout.column(heading="Upscale", align=True) + col = col.column(heading="Upscale", align=True) col.prop(context.scene, "video_to_video", text="2x") if type == "image" and ( image_model_card == "stabilityai/stable-diffusion-xl-base-1.0" ): - col = layout.column(heading="Refine", align=True) + col = col.column(heading="Refine", align=True) col.prop(context.scene, "refine_sd", text="Image") sub_col = col.row() sub_col.active = context.scene.refine_sd - col = layout.column() - col.prop(context.scene, "input_strips", text="Input") - if input == "input_strips": - col.prop(context.scene, "image_power", text="Strip Power") + col.prop(context.scene, "movie_num_batch", text="Batch Count") col = layout.column() - col.prop(context.scene, "generatorai_typeselect", text="Output") - col.prop(context.scene, "movie_num_batch", text="Batch Count") + col = col.box() if input == "input_strips": - row = layout.row(align=True) - row.scale_y = 1.1 + ed = scene.sequence_editor + + row = col.row(align=True) + row.scale_y = 1.2 row.operator("sequencer.text_to_generator", text="Generate from Strips") else: - row = layout.row(align=True) - row.scale_y = 1.1 + row = col.row(align=True) + row.scale_y = 1.2 if type == "movie": if movie_model_card == "stabilityai/stable-diffusion-xl-base-1.0": row.operator("sequencer.text_to_generator", text="Generate from Strips") @@ -1098,12 +1122,13 @@ class SEQUENCER_OT_generate_movie(Operator): if (scene.movie_path or scene.image_path) and input == "input_strips": if movie_model_card == "stabilityai/stable-diffusion-xl-base-1.0": #img2img - from diffusers import StableDiffusionXLImg2ImgPipeline - + from diffusers import StableDiffusionXLImg2ImgPipeline, AutoencoderKL + vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16, force_upcast=False) pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained( movie_model_card, torch_dtype=torch.float16, variant="fp16", + vae=vae, ) from diffusers import DPMSolverMultistepScheduler @@ -1130,25 +1155,22 @@ class SEQUENCER_OT_generate_movie(Operator): ) if low_vram: - #refiner.unet.enable_forward_chunking(chunk_size=1, dim=1) #Heavy + refiner.enable_model_cpu_offload() refiner.enable_vae_slicing() else: refiner.to("cuda") - else: + else: #vid2vid if movie_model_card == "cerspense/zeroscope_v2_dark_30x448x256" or movie_model_card == "cerspense/zeroscope_v2_576w" or scene.image_path: card = "cerspense/zeroscope_v2_XL" - safe = False else: card = movie_model_card - safe = True from diffusers import VideoToVideoSDPipeline - upscale = VideoToVideoSDPipeline.from_pretrained( card, torch_dtype=torch.float16, - use_safetensors=safe, + #use_safetensors=True, ) from diffusers import DPMSolverMultistepScheduler @@ -1156,9 +1178,9 @@ class SEQUENCER_OT_generate_movie(Operator): upscale.scheduler = DPMSolverMultistepScheduler.from_config(upscale.scheduler.config) if low_vram: - torch.cuda.set_per_process_memory_fraction(0.98) + #torch.cuda.set_per_process_memory_fraction(0.98) upscale.enable_model_cpu_offload() - upscale.unet.enable_forward_chunking(chunk_size=1, dim=1) # here: + upscale.unet.enable_forward_chunking(chunk_size=1, dim=1) # heavy: upscale.enable_vae_slicing() else: upscale.to("cuda") @@ -1166,6 +1188,7 @@ class SEQUENCER_OT_generate_movie(Operator): # Models for movie generation else: from diffusers import TextToVideoSDPipeline + pipe = TextToVideoSDPipeline.from_pretrained( movie_model_card, torch_dtype=torch.float16, @@ -1615,6 +1638,23 @@ class SEQUENCER_OT_generate_audio(Operator): return {"FINISHED"} +def find_strip_by_name(scene, name): + for sequence in scene.sequence_editor.sequences: + if sequence.name == name: + return sequence + return None + + +def get_strip_path(strip): + if strip: + strip_dirname = os.path.dirname(strip.directory) + image_path = bpy.path.abspath( + os.path.join(strip_dirname, strip.elements[0].filename) + ) + return image_path + return None + + class SEQUENCER_OT_generate_image(Operator): """Generate Image""" @@ -1655,6 +1695,8 @@ class SEQUENCER_OT_generate_image(Operator): torch.cuda.empty_cache() current_frame = scene.frame_current + type = scene.generatorai_typeselect + input = scene.input_strips prompt = style_prompt(scene.generate_movie_prompt)[0] negative_prompt = scene.generate_movie_negative_prompt +", "+ style_prompt(scene.generate_movie_prompt)[1] +", nsfw nude nudity" image_x = scene.generate_movie_x @@ -1668,16 +1710,54 @@ class SEQUENCER_OT_generate_image(Operator): preferences = context.preferences addon_prefs = preferences.addons[__name__].preferences image_model_card = addon_prefs.image_model_card - do_refine = (scene.refine_sd or scene.image_path) and image_model_card == "stabilityai/stable-diffusion-xl-base-1.0" + do_inpaint = (input == "input_strips" and type == "image" and scene.inpaint_selected_strip) + do_refine = False #(scene.refine_sd or scene.image_path) and image_model_card == "stabilityai/stable-diffusion-xl-base-1.0" # LOADING MODELS print("\nModel: " + image_model_card) - # Models for stable diffusion - if not image_model_card == "DeepFloyd/IF-I-M-v1.0": + # models for inpaint + if do_inpaint: + + from diffusers import StableDiffusionXLInpaintPipeline, AutoencoderKL + from diffusers.utils import load_image + + # clear the VRAM + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16, force_upcast=False) + pipe = StableDiffusionXLInpaintPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", vae=vae, torch_dtype=torch.float16, variant="fp16", force_upcast=False) #use_safetensors=True + if low_vram: + #torch.cuda.set_per_process_memory_fraction(0.99) + pipe.enable_model_cpu_offload() + pipe.enable_vae_slicing() + else: + pipe.to("cuda") + +# refiner = StableDiffusionXLInpaintPipeline.from_pretrained( +# "stabilityai/stable-diffusion-xl-refiner-1.0", +# text_encoder_2=pipe.text_encoder_2, +# vae = vae, +# #vae=pipe.vae, +# torch_dtype=torch.float16, +# use_safetensors=True, +# variant="fp16", +# ) +# if low_vram: +# refiner.enable_model_cpu_offload() +# refiner.enable_vae_slicing() +# else: +# refiner.to("cuda") + + # Models for stable diffusion + elif not image_model_card == "DeepFloyd/IF-I-M-v1.0": + from diffusers import AutoencoderKL + vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16, force_upcast=False) pipe = DiffusionPipeline.from_pretrained( image_model_card, + vae=vae, torch_dtype=torch.float16, variant="fp16", ) @@ -1747,11 +1827,12 @@ class SEQUENCER_OT_generate_image(Operator): # Add refiner model if chosen. if do_refine: print("Refine Model: " + "stabilityai/stable-diffusion-xl-refiner-1.0") - from diffusers import StableDiffusionXLImg2ImgPipeline + from diffusers import StableDiffusionXLImg2ImgPipeline, AutoencoderKL + vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16, force_upcast=False) refiner = StableDiffusionXLImg2ImgPipeline.from_pretrained( "stabilityai/stable-diffusion-xl-refiner-1.0", text_encoder_2=pipe.text_encoder_2, - vae=pipe.vae, + vae=vae, torch_dtype=torch.float16, variant="fp16", ) @@ -1836,6 +1917,28 @@ class SEQUENCER_OT_generate_image(Operator): # image[0].save("./if_stage_III.png") image = image[0] + # Inpaint + elif do_inpaint: + print("Process: Inpaint") + img_path = load_image(scene.image_path).convert("RGB") + + mask_strip =find_strip_by_name(scene, scene.inpaint_selected_strip) + if not mask_strip: + return + + mask_path = get_strip_path(mask_strip) + + init_image = load_image(img_path).convert("RGB") + mask_image = load_image(mask_path).convert("RGB") + + image = pipe( + prompt=prompt, + image=init_image, + mask_image=mask_image, + num_inference_steps=image_num_inference_steps, + strength=1.00 - scene.image_power, + ).images[0] + # Img2img elif scene.image_path: print("Process: Image to image") @@ -2003,7 +2106,7 @@ class SEQUENCER_OT_strip_to_generatorAI(Operator): if len(strips) > 1: styled_prompt = style_prompt(strip_prompt + ", " + prompt)[0] - styled_prompt = style_prompt(strip_prompt + ", " + prompt)[1] + styled_negative_prompt = style_prompt(strip_prompt + ", " + prompt)[1] else: styled_prompt = style_prompt(prompt)[0] styled_negative_prompt = style_prompt(prompt)[1] @@ -2048,7 +2151,7 @@ class SEQUENCER_OT_strip_to_generatorAI(Operator): if len(strips) > 1: styled_prompt = style_prompt(strip_prompt + ", " + prompt)[0] - styled_prompt = style_prompt(strip_prompt + ", " + prompt)[1] + styled_negative_prompt = style_prompt(strip_prompt + ", " + prompt)[1] else: styled_prompt = style_prompt(prompt)[0] styled_negative_prompt = style_prompt(prompt)[1] @@ -2227,6 +2330,9 @@ def register(): default="en", ) + # Inpaint + bpy.types.Scene.inpaint_selected_strip = bpy.props.StringProperty(name="inpaint_selected_strip", default="") + # Upscale bpy.types.Scene.video_to_video = bpy.props.BoolProperty( name="video_to_video", @@ -2271,6 +2377,7 @@ def register(): default="no_style", ) + for cls in classes: bpy.utils.register_class(cls) @@ -2292,6 +2399,7 @@ def unregister(): del bpy.types.Scene.image_path del bpy.types.Scene.refine_sd del bpy.types.Scene.generatorai_styles + del bpy.types.Scene.inpaint_selected_strip if __name__ == "__main__": unregister()