add stable-video-diffusion

modules/processing_diffusers.py

@@ -5,7 +5,6 @@ import inspect
 import typing
 import torch
 import torchvision.transforms.functional as TF
-import diffusers
 import modules.devices as devices
 import modules.shared as shared
 import modules.sd_samplers as sd_samplers
@@ -73,10 +72,7 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro
                 time.sleep(0.1)
 
     def diffusers_callback(_pipe, step: int, timestep: int, kwargs: dict):
-        latents = correction_callback(p, timestep, kwargs)
-        latents = kwargs['latents']
         shared.state.sampling_step = step
-        shared.state.current_latent = latents
         if shared.state.interrupted or shared.state.skipped:
             raise AssertionError('Interrupted...')
         if shared.state.paused:
@@ -85,7 +81,12 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro
                 if shared.state.interrupted or shared.state.skipped:
                     raise AssertionError('Interrupted...')
                 time.sleep(0.1)
-        return {'latents': latents}
+        if kwargs.get('latents', None) is None:
+            print('HERE NO')
+            return kwargs
+        kwargs = correction_callback(p, timestep, kwargs)
+        shared.state.current_latent = kwargs['latents']
+        return kwargs
 
     def full_vae_decode(latents, model):
         t0 = time.time()
@@ -129,7 +130,7 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro
         if len(latents) == 0:
             return []
         decoded = torch.zeros((len(latents), 3, latents.shape[2] * 8, latents.shape[3] * 8), dtype=devices.dtype_vae, device=devices.device)
-        for i in range(len(output.images)):
+        for i in range(latents.shape[0]):
             decoded[i] = sd_vae_taesd.decode(latents[i])
         return decoded
 
@@ -151,6 +152,8 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro
         if not hasattr(model, 'vae'):
             shared.log.error('VAE not found in model')
             return []
+        if latents.shape[0] == 4 and latents.shape[1] != 4: # likely animatediff latent
+            latents = latents.permute(1, 0, 2, 3)
         if len(latents.shape) == 3: # lost a batch dim in hires
             latents = latents.unsqueeze(0)
         if full_quality:
@@ -200,16 +203,27 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro
 
     def task_specific_kwargs(model):
         task_args = {}
-        is_img2img_model = bool("Zero123" in shared.sd_model.__class__.__name__)
+        is_img2img_model = bool('Zero123' in shared.sd_model.__class__.__name__)
         if sd_models.get_diffusers_task(model) == sd_models.DiffusersTaskType.TEXT_2_IMAGE and not is_img2img_model:
             p.ops.append('txt2img')
-            task_args = {"height": 8 * math.ceil(p.height / 8), "width": 8 * math.ceil(p.width / 8)}
+            task_args = {
+                'height': 8 * math.ceil(p.height / 8),
+                'width': 8 * math.ceil(p.width / 8),
+            }
         elif (sd_models.get_diffusers_task(model) == sd_models.DiffusersTaskType.IMAGE_2_IMAGE or is_img2img_model) and len(getattr(p, 'init_images' ,[])) > 0:
             p.ops.append('img2img')
-            task_args = {"image": p.init_images, "strength": p.denoising_strength}
+            task_args = {
+                'image': p.init_images,
+                'strength': p.denoising_strength,
+            }
         elif sd_models.get_diffusers_task(model) == sd_models.DiffusersTaskType.INSTRUCT and len(getattr(p, 'init_images' ,[])) > 0:
             p.ops.append('instruct')
-            task_args = {"height": 8 * math.ceil(p.height / 8), "width": 8 * math.ceil(p.width / 8), "image": p.init_images, "strength": p.denoising_strength}
+            task_args = {
+                'height': 8 * math.ceil(p.height / 8),
+                'width': 8 * math.ceil(p.width / 8),
+                'image': p.init_images,
+                'strength': p.denoising_strength,
+            }
         elif (sd_models.get_diffusers_task(model) == sd_models.DiffusersTaskType.INPAINTING or is_img2img_model) and len(getattr(p, 'init_images' ,[])) > 0:
             p.ops.append('inpaint')
             if getattr(p, 'mask', None) is None:
@@ -217,7 +231,13 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro
             width = 8 * math.ceil(p.init_images[0].width / 8)
             height = 8 * math.ceil(p.init_images[0].height / 8)
             # option-1: use images as inputs
-            task_args = {"image": p.init_images, "mask_image": p.mask, "strength": p.denoising_strength, "height": height, "width": width}
+            task_args = {
+                'image': p.init_images,
+                'mask_image': p.mask,
+                'strength': p.denoising_strength,
+                'height': height,
+                'width': width,
+            }
             """ # option-2: preprocess images into latents using diffusers
             vae_scale_factor = 2 ** (len(model.vae.config.block_out_channels) - 1)
             image_processor = diffusers.image_processor.VaeImageProcessor(vae_scale_factor=vae_scale_factor)
@@ -237,11 +257,16 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro
             task_args = {"image": p.init_images, "mask_image": mask_image, "masked_image_latents": masked_image_latents, "strength": p.denoising_strength, "height": height, "width": width}
             """
         if model.__class__.__name__ == 'LatentConsistencyModelPipeline' and hasattr(p, 'init_images') and len(p.init_images) > 0:
+            p.ops.append('lcm')
             init_latents = [vae_encode(image, model=shared.sd_model, full_quality=p.full_quality).squeeze(dim=0) for image in p.init_images]
             init_latent = torch.stack(init_latents, dim=0).to(shared.device)
             init_noise = p.denoising_strength * create_random_tensors(init_latent.shape[1:], seeds=p.all_seeds, subseeds=p.all_subseeds, subseed_strength=p.subseed_strength, p=p)
             init_latent = (1 - p.denoising_strength) * init_latent + init_noise
-            task_args = {"latents": init_latent.to(model.dtype), "width": p.width, "height": p.height }
+            task_args = {
+                'latents': init_latent.to(model.dtype),
+                'width': p.width,
+                'height': p.height,
+            }
         return task_args
 
     def set_pipeline_args(model, prompts: list, negative_prompts: list, prompts_2: typing.Optional[list]=None, negative_prompts_2: typing.Optional[list]=None, desc:str='', **kwargs):
@@ -309,12 +334,13 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro
                 pass
         task_kwargs = task_specific_kwargs(model)
         for arg in task_kwargs:
-            if arg in possible and arg not in args: # task specific args should not override args
+            # if arg in possible and arg not in args: # task specific args should not override args
+            if arg in possible:
                 args[arg] = task_kwargs[arg]
-            else:
-                pass
-        for k, v in getattr(p, 'task_args', {}).items():
+        task_args = getattr(p, 'task_args', {})
+        for k, v in task_args.items():
             args[k] = v
+
         hypertile_set(p, hr=len(getattr(p, 'init_images', [])))
         clean = args.copy()
         clean.pop('callback', None)
@@ -394,10 +420,6 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro
             # TODO extra_generation_params add sampler options
             # p.extra_generation_params['Sampler options'] = ''
 
-    recompile_model()
-    update_sampler(shared.sd_model)
-    p.extra_generation_params['Pipeline'] = shared.sd_model.__class__.__name__
-
     if len(getattr(p, 'init_images', [])) > 0:
         while len(p.init_images) < len(prompts):
             p.init_images.append(p.init_images[-1])
@@ -475,7 +497,10 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro
         clip_skip=p.clip_skip,
         desc='Base',
     )
+    recompile_model()
+    update_sampler(shared.sd_model)
     shared.state.sampling_steps = base_args['num_inference_steps']
+    p.extra_generation_params['Pipeline'] = shared.sd_model.__class__.__name__
     p.extra_generation_params["Sampler Eta"] = shared.opts.scheduler_eta if shared.opts.scheduler_eta is not None and shared.opts.scheduler_eta > 0 and shared.opts.scheduler_eta < 1 else None
     try:
         output = shared.sd_model(**base_args) # pylint: disable=not-callable
@@ -486,9 +511,13 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro
         shared.log.info(e)
     except ValueError as e:
         shared.state.interrupted = True
-        shared.log.error(f'Processing: {e}')
+        shared.log.error(f'Processing: args={base_args} {e}')
         if shared.cmd_opts.debug:
             errors.display(e, 'Processing')
+    except RuntimeError as e:
+        shared.state.interrupted = True
+        shared.log.error(f'Processing: args={base_args} {e}')
+        errors.display(e, 'Processing')
 
     if hasattr(shared.sd_model, 'embedding_db') and len(shared.sd_model.embedding_db.embeddings_used) > 0:
         p.extra_generation_params['Embeddings'] = ', '.join(shared.sd_model.embedding_db.embeddings_used)