# # For licensing see accompanying LICENSE.md file. # Copyright (C) 2022 Apple Inc. All Rights Reserved. # import argparse from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline from diffusers.pipelines.pipeline_utils import DiffusionPipeline from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput from diffusers.schedulers import ( DDIMScheduler, DPMSolverMultistepScheduler, EulerAncestralDiscreteScheduler, EulerDiscreteScheduler, LMSDiscreteScheduler, PNDMScheduler, ) from diffusers.schedulers.scheduling_utils import SchedulerMixin import gc import inspect import logging logging.basicConfig() logger = logging.getLogger(__name__) logger.setLevel(logging.INFO) import numpy as np import os from python_coreml_stable_diffusion.coreml_model import ( CoreMLModel, _load_mlpackage, _load_mlpackage_controlnet, get_available_compute_units, ) import time import torch # Only used for `torch.from_tensor` in `pipe.scheduler.step()` from transformers import CLIPFeatureExtractor, CLIPTokenizer from typing import List, Optional, Union, Tuple from PIL import Image class CoreMLStableDiffusionPipeline(DiffusionPipeline): """ Core ML version of `diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline` """ def __init__( self, text_encoder: CoreMLModel, unet: CoreMLModel, vae_decoder: CoreMLModel, scheduler: Union[ DDIMScheduler, DPMSolverMultistepScheduler, EulerAncestralDiscreteScheduler, EulerDiscreteScheduler, LMSDiscreteScheduler, PNDMScheduler ], tokenizer: CLIPTokenizer, controlnet: Optional[List[CoreMLModel]], xl: Optional[bool] = False, force_zeros_for_empty_prompt: Optional[bool] = True, feature_extractor: Optional[CLIPFeatureExtractor] = None, safety_checker: Optional[CoreMLModel] = None, text_encoder_2: Optional[CoreMLModel] = None, tokenizer_2: Optional[CLIPTokenizer] = None ): super().__init__() # Register non-Core ML components of the pipeline similar to the original pipeline self.register_modules( tokenizer=tokenizer, scheduler=scheduler, feature_extractor=feature_extractor, ) if safety_checker is None: # Reproduce original warning: # https://github.com/huggingface/diffusers/blob/v0.9.0/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py#L119 logger.warning( f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure" " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered" " results in services or applications open to the public. Both the diffusers team and Hugging Face" " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling" " it only for use-cases that involve analyzing network behavior or auditing its results. For more" " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ." ) self.xl = xl self.force_zeros_for_empty_prompt = force_zeros_for_empty_prompt # Register Core ML components of the pipeline self.safety_checker = safety_checker self.text_encoder = text_encoder self.text_encoder_2 = text_encoder_2 self.tokenizer_2 = tokenizer_2 self.unet = unet self.unet.in_channels = self.unet.expected_inputs["sample"]["shape"][1] self.controlnet = controlnet self.vae_decoder = vae_decoder VAE_DECODER_UPSAMPLE_FACTOR = 8 # In PyTorch, users can determine the tensor shapes dynamically by default # In CoreML, tensors have static shapes unless flexible shapes were used during export # See https://coremltools.readme.io/docs/flexible-inputs latent_h, latent_w = self.unet.expected_inputs["sample"]["shape"][2:] self.height = latent_h * VAE_DECODER_UPSAMPLE_FACTOR self.width = latent_w * VAE_DECODER_UPSAMPLE_FACTOR logger.info( f"Stable Diffusion configured to generate {self.height}x{self.width} images" ) def _encode_prompt(self, prompt, prompt_2: Optional[str] = None, do_classifier_free_guidance: bool = True, negative_prompt: Optional[str] = None, negative_prompt_2: Optional[str] = None, ): batch_size = len(prompt) if isinstance(prompt, list) else 1 if self.xl is True: prompts = [prompt, prompt_2] if prompt_2 is not None else [prompt, prompt] # refiner uses only one tokenizer and text encoder (tokenizer_2 and text_encoder_2) tokenizers = [self.tokenizer, self.tokenizer_2] if self.tokenizer is not None else [self.tokenizer_2] text_encoders = [self.text_encoder, self.text_encoder_2] if self.text_encoder is not None else [ self.text_encoder_2] hidden_state_key = 'hidden_embeds' else: prompts = [prompt] tokenizers = [self.tokenizer] text_encoders = [self.text_encoder] hidden_state_key = 'last_hidden_state' prompt_embeds_list = [] for prompt, tokenizer, text_encoder in zip(prompts, tokenizers, text_encoders): text_inputs = tokenizer( prompt, padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="np", ) text_input_ids = text_inputs.input_ids # tokenize without max_length to catch any truncation untruncated_ids = tokenizer(prompt, padding="longest", return_tensors="np").input_ids if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not np.equal( text_input_ids, untruncated_ids ): removed_text = tokenizer.batch_decode(untruncated_ids[:, tokenizer.model_max_length - 1: -1]) logger.warning( "The following part of your input was truncated because CLIP can only handle sequences up to" f" {tokenizer.model_max_length} tokens: {removed_text}" ) embeddings = text_encoder(input_ids=text_input_ids.astype(np.float32)) prompt_embeds_list.append(embeddings[hidden_state_key]) # We are only ALWAYS interested in the pooled output of the final text encoder if self.xl: pooled_prompt_embeds = embeddings['pooled_outputs'] prompt_embeds = np.concatenate(prompt_embeds_list, axis=-1) if do_classifier_free_guidance and negative_prompt is None and self.force_zeros_for_empty_prompt: negative_prompt_embeds = np.zeros_like(prompt_embeds) if self.xl: negative_pooled_prompt_embeds = np.zeros_like(pooled_prompt_embeds) elif do_classifier_free_guidance: negative_prompt = negative_prompt or "" negative_prompt_2 = negative_prompt_2 or negative_prompt # normalize str to list negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt negative_prompt_2 = ( batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2 ) uncond_tokens: List[str] if prompts is not None and type(prompts) is not type(negative_prompt): raise TypeError( f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" f" {type(prompt)}." ) elif batch_size != len(negative_prompt): raise ValueError( f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:" f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches" " the batch size of `prompt`.") else: uncond_tokens = [negative_prompt, negative_prompt_2] negative_prompt_embeds_list = [] for negative_prompt, tokenizer, text_encoder in zip(uncond_tokens, tokenizers, text_encoders): max_length = prompt_embeds.shape[1] uncond_input = tokenizer( negative_prompt, padding="max_length", max_length=max_length, truncation=True, return_tensors="np", ) uncond_input_ids = uncond_input.input_ids negative_embeddings = text_encoder( input_ids=uncond_input_ids.astype(np.float32) ) negative_text_embeddings = negative_embeddings[hidden_state_key] negative_prompt_embeds_list.append(negative_text_embeddings) # We are only ALWAYS interested in the pooled output of the final text encoder if self.xl: negative_pooled_prompt_embeds = negative_embeddings['pooled_outputs'] negative_prompt_embeds = np.concatenate(negative_prompt_embeds_list, axis=-1) if do_classifier_free_guidance: # For classifier free guidance, we need to do two forward passes. # Here we concatenate the unconditional and text embeddings into a single batch # to avoid doing two forward passes prompt_embeds = np.concatenate( [negative_prompt_embeds, prompt_embeds]) if self.xl: pooled_prompt_embeds = np.concatenate( [negative_pooled_prompt_embeds, pooled_prompt_embeds]) prompt_embeddings = prompt_embeds.transpose(0, 2, 1)[:, :, None, :] if self.xl: return prompt_embeddings, pooled_prompt_embeds else: return prompt_embeddings, None def run_controlnet(self, sample, timestep, encoder_hidden_states, controlnet_cond, output_dtype=np.float16): if not self.controlnet: raise ValueError( "Conditions for controlnet are given but the pipeline has no controlnet modules") for i, (module, cond) in enumerate(zip(self.controlnet, controlnet_cond)): module_outputs = module( sample=sample.astype(np.float16), timestep=timestep.astype(np.float16), encoder_hidden_states=encoder_hidden_states.astype(np.float16), controlnet_cond=cond.astype(np.float16), ) if i == 0: outputs = module_outputs else: for key in outputs.keys(): outputs[key] += module_outputs[key] outputs = {k: v.astype(output_dtype) for k, v in outputs.items()} return outputs def run_safety_checker(self, image): if self.safety_checker is not None: safety_checker_input = self.feature_extractor( self.numpy_to_pil(image), return_tensors="np", ) safety_checker_outputs = self.safety_checker( clip_input=safety_checker_input.pixel_values.astype( np.float16), images=image.astype(np.float16), adjustment=np.array([0.]).astype( np.float16), # defaults to 0 in original pipeline ) # Unpack dict has_nsfw_concept = safety_checker_outputs["has_nsfw_concepts"] image = safety_checker_outputs["filtered_images"] concept_scores = safety_checker_outputs["concept_scores"] logger.info( f"Generated image has nsfw concept={has_nsfw_concept.any()}") else: has_nsfw_concept = None return image, has_nsfw_concept def decode_latents(self, latents): latents = 1 / 0.18215 * latents dtype = self.vae_decoder.expected_inputs['z']['dtype'] image = self.vae_decoder(z=latents.astype(dtype))["image"] image = np.clip(image / 2 + 0.5, 0, 1) image = image.transpose((0, 2, 3, 1)) return image def prepare_latents(self, batch_size, num_channels_latents, height, width, latents=None): latents_shape = (batch_size, num_channels_latents, self.height // 8, self.width // 8) if latents is None: latents = np.random.randn(*latents_shape).astype(np.float16) elif latents.shape != latents_shape: raise ValueError( f"Unexpected latents shape, got {latents.shape}, expected {latents_shape}" ) init_noise = self.scheduler.init_noise_sigma if isinstance(init_noise, torch.Tensor): init_noise = init_noise.numpy() latents = latents * init_noise return latents def prepare_control_cond(self, controlnet_cond, do_classifier_free_guidance, batch_size, num_images_per_prompt): processed_cond_list = [] for cond in controlnet_cond: cond = np.stack([cond] * batch_size * num_images_per_prompt) if do_classifier_free_guidance: cond = np.concatenate([cond] * 2) processed_cond_list.append(cond) return processed_cond_list def check_inputs(self, prompt, height, width, callback_steps): if height != self.height or width != self.width: logger.warning( "`height` and `width` dimensions (of the output image tensor) are fixed when exporting the Core ML models " \ "unless flexible shapes are used during export (https://coremltools.readme.io/docs/flexible-inputs). " \ "This pipeline was provided with Core ML models that generate {self.height}x{self.width} images (user requested {height}x{width})" ) if not isinstance(prompt, str) and not isinstance(prompt, list): raise ValueError( f"`prompt` has to be of type `str` or `list` but is {type(prompt)}" ) if height % 8 != 0 or width % 8 != 0: raise ValueError( f"`height` and `width` have to be divisible by 8 but are {height} and {width}." ) if (callback_steps is None) or (callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0)): raise ValueError( f"`callback_steps` has to be a positive integer but is {callback_steps} of type" f" {type(callback_steps)}.") def prepare_extra_step_kwargs(self, eta): # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502 # and should be between [0, 1] accepts_eta = "eta" in set( inspect.signature(self.scheduler.step).parameters.keys()) extra_step_kwargs = {} if accepts_eta: extra_step_kwargs["eta"] = eta return extra_step_kwargs def _get_add_time_ids(self, original_size, crops_coords_top_left, target_size, dtype): add_time_ids = list(original_size + crops_coords_top_left + target_size) add_time_ids = np.array(add_time_ids).astype(dtype) return add_time_ids def __call__( self, prompt, height=512, width=512, num_inference_steps=50, guidance_scale=7.5, negative_prompt=None, num_images_per_prompt=1, eta=0.0, latents=None, output_type="pil", return_dict=True, callback=None, callback_steps=1, controlnet_cond=None, original_size: Optional[Tuple[int, int]]=None, crops_coords_top_left: Tuple[int, int]=(0, 0), target_size: Optional[Tuple[int, int]]=None, unet_batch_one=False, **kwargs, ): # 1. Check inputs. Raise error if not correct self.check_inputs(prompt, height, width, callback_steps) height = height or self.height width = width or self.width original_size = original_size or (height, width) target_size = target_size or (height, width) # 2. Define call parameters batch_size = 1 if isinstance(prompt, str) else len(prompt) if batch_size > 1 or num_images_per_prompt > 1: raise NotImplementedError( "For batched generation of multiple images and/or multiple prompts, please refer to the Swift package." ) # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2) # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1` # corresponds to doing no classifier free guidance. do_classifier_free_guidance = guidance_scale > 1.0 # 3. Encode input prompt text_embeddings, pooled_prompt_embeds = self._encode_prompt( prompt=prompt, prompt_2=None, do_classifier_free_guidance=do_classifier_free_guidance, negative_prompt=negative_prompt, negative_prompt_2=None ) # 4. Prepare XL kwargs if needed unet_additional_kwargs = {} # we add pooled prompt embeds + time_ids to unet kwargs if self.xl: add_text_embeds = pooled_prompt_embeds add_time_ids = self._get_add_time_ids(original_size, crops_coords_top_left, target_size, text_embeddings.dtype) if do_classifier_free_guidance: # TODO: This checks if the time_ids input is looking for time_ids.shape == (12,) or (2, 6) # Remove once model input shapes are ubiquitous if len(self.unet.expected_inputs['time_ids']['shape']) > 1: add_time_ids = [add_time_ids] add_time_ids = np.concatenate([add_time_ids, add_time_ids]) unet_additional_kwargs.update({'text_embeds': add_text_embeds.astype(np.float16), 'time_ids': add_time_ids.astype(np.float16)}) # 5. Prepare timesteps self.scheduler.set_timesteps(num_inference_steps) timesteps = self.scheduler.timesteps # 6. Prepare latent variables and controlnet cond num_channels_latents = self.unet.in_channels latents = self.prepare_latents( batch_size * num_images_per_prompt, num_channels_latents, height, width, latents, ) if controlnet_cond: controlnet_cond = self.prepare_control_cond( controlnet_cond, do_classifier_free_guidance, batch_size, num_images_per_prompt, ) # 7. Prepare extra step kwargs extra_step_kwargs = self.prepare_extra_step_kwargs(eta) # 8. Denoising loop for i, t in enumerate(self.progress_bar(timesteps)): # expand the latents if we are doing classifier free guidance latent_model_input = np.concatenate( [latents] * 2) if do_classifier_free_guidance else latents latent_model_input = self.scheduler.scale_model_input( latent_model_input, t) if isinstance(latent_model_input, torch.Tensor): latent_model_input = latent_model_input.numpy() if do_classifier_free_guidance: timestep = np.array([t, t], np.float16) else: timestep = np.array([t,], np.float16) # controlnet if controlnet_cond: control_net_additional_residuals = self.run_controlnet( sample=latent_model_input, timestep=timestep, encoder_hidden_states=text_embeddings, controlnet_cond=controlnet_cond, ) else: control_net_additional_residuals = {} # predict the noise residual unet_additional_kwargs.update(control_net_additional_residuals) # get prediction from unet if not (unet_batch_one and do_classifier_free_guidance): noise_pred = self.unet( sample=latent_model_input.astype(np.float16), timestep=timestep, encoder_hidden_states=text_embeddings.astype(np.float16), **unet_additional_kwargs, )["noise_pred"] if do_classifier_free_guidance: noise_pred_uncond, noise_pred_text = np.split(noise_pred, 2) else: # query unet sequentially latent_model_input = latent_model_input.astype(np.float16) text_embeddings = text_embeddings.astype(np.float16) timestep = np.array([t,], np.float16) noise_pred_uncond = self.unet( sample=np.expand_dims(latent_model_input[0], axis=0), timestep=timestep, encoder_hidden_states=np.expand_dims(text_embeddings[0], axis=0), **unet_additional_kwargs, )["noise_pred"] noise_pred_text = self.unet( sample=np.expand_dims(latent_model_input[1], axis=0), timestep=timestep, encoder_hidden_states=np.expand_dims(text_embeddings[1], axis=0), **unet_additional_kwargs, )["noise_pred"] # perform guidance if do_classifier_free_guidance: noise_pred = noise_pred_uncond + guidance_scale * ( noise_pred_text - noise_pred_uncond) # compute the previous noisy sample x_t -> x_t-1 latents = self.scheduler.step(torch.from_numpy(noise_pred), t, torch.from_numpy(latents), **extra_step_kwargs, ).prev_sample.numpy() # call the callback, if provided if callback is not None and i % callback_steps == 0: callback(i, t, latents) # 8. Post-processing image = self.decode_latents(latents) # 9. Run safety checker image, has_nsfw_concept = self.run_safety_checker(image) # 10. Convert to PIL if output_type == "pil": image = self.numpy_to_pil(image) if not return_dict: return (image, has_nsfw_concept) return StableDiffusionPipelineOutput( images=image, nsfw_content_detected=has_nsfw_concept) def get_available_schedulers(): schedulers = {} for scheduler in [DDIMScheduler, DPMSolverMultistepScheduler, EulerAncestralDiscreteScheduler, EulerDiscreteScheduler, LMSDiscreteScheduler, PNDMScheduler]: schedulers[scheduler().__class__.__name__.replace("Scheduler", "")] = scheduler return schedulers SCHEDULER_MAP = get_available_schedulers() def get_coreml_pipe(pytorch_pipe, mlpackages_dir, model_version, compute_unit, delete_original_pipe=True, scheduler_override=None, controlnet_models=None, force_zeros_for_empty_prompt=True, sources=None): """ Initializes and returns a `CoreMLStableDiffusionPipeline` from an original diffusers PyTorch pipeline sources: 'packages' or 'compiled' forces creation of model from specified sources. sources must be in mlpackages_dir """ # Ensure `scheduler_override` object is of correct type if specified if scheduler_override is not None: assert isinstance(scheduler_override, SchedulerMixin) logger.warning( "Overriding scheduler in pipeline: " f"Default={pytorch_pipe.scheduler}, Override={scheduler_override}") # Gather configured tokenizer and scheduler attributes from the original pipe if 'xl' in model_version: coreml_pipe_kwargs = { "tokenizer": pytorch_pipe.tokenizer, 'tokenizer_2': pytorch_pipe.tokenizer_2, "scheduler": pytorch_pipe.scheduler if scheduler_override is None else scheduler_override, 'xl': True, } model_packages_to_load = ["text_encoder", "text_encoder_2", "unet", "vae_decoder"] else: coreml_pipe_kwargs = { "tokenizer": pytorch_pipe.tokenizer, "scheduler": pytorch_pipe.scheduler if scheduler_override is None else scheduler_override, "feature_extractor": pytorch_pipe.feature_extractor, } model_packages_to_load = ["text_encoder", "unet", "vae_decoder"] coreml_pipe_kwargs["force_zeros_for_empty_prompt"] = force_zeros_for_empty_prompt if getattr(pytorch_pipe, "safety_checker", None) is not None: model_packages_to_load.append("safety_checker") else: logger.warning( f"Original diffusers pipeline for {model_version} does not have a safety_checker, " "Core ML pipeline will mirror this behavior.") coreml_pipe_kwargs["safety_checker"] = None if delete_original_pipe: del pytorch_pipe gc.collect() logger.info("Removed PyTorch pipe to reduce peak memory consumption") if controlnet_models: model_packages_to_load.remove("unet") coreml_pipe_kwargs["unet"] = _load_mlpackage( submodule_name="control-unet", mlpackages_dir=mlpackages_dir, model_version=model_version, compute_unit=compute_unit, ) coreml_pipe_kwargs["controlnet"] = [_load_mlpackage_controlnet( mlpackages_dir, model_version, compute_unit, ) for model_version in controlnet_models] else: coreml_pipe_kwargs["controlnet"] = None # Load Core ML models logger.info(f"Loading Core ML models in memory from {mlpackages_dir}") coreml_pipe_kwargs.update({ model_name: _load_mlpackage( submodule_name=model_name, mlpackages_dir=mlpackages_dir, model_version=model_version, compute_unit=compute_unit, sources=sources, ) for model_name in model_packages_to_load }) logger.info("Done.") logger.info("Initializing Core ML pipe for image generation") coreml_pipe = CoreMLStableDiffusionPipeline(**coreml_pipe_kwargs) logger.info("Done.") return coreml_pipe def get_image_path(args, **override_kwargs): """ mkdir output folder and encode metadata in the filename """ out_folder = os.path.join(args.o, "_".join(args.prompt.replace("/", "_").rsplit(" "))) os.makedirs(out_folder, exist_ok=True) out_fname = f"randomSeed_{override_kwargs.get('seed', None) or args.seed}" out_fname += f"_computeUnit_{override_kwargs.get('compute_unit', None) or args.compute_unit}" out_fname += f"_modelVersion_{override_kwargs.get('model_version', None) or args.model_version.replace('/', '_')}" if args.scheduler is not None: out_fname += f"_customScheduler_{override_kwargs.get('scheduler', None) or args.scheduler}" out_fname += f"_numInferenceSteps{override_kwargs.get('num_inference_steps', None) or args.num_inference_steps}" return os.path.join(out_folder, out_fname + ".png") def prepare_controlnet_cond(image_path, height, width): image = Image.open(image_path).convert("RGB") image = image.resize((height, width), resample=Image.LANCZOS) image = np.array(image).transpose(2, 0, 1) / 255.0 return image def main(args): logger.info(f"Setting random seed to {args.seed}") np.random.seed(args.seed) logger.info("Initializing PyTorch pipe for reference configuration") SDP = StableDiffusionXLPipeline if 'xl' in args.model_version else StableDiffusionPipeline pytorch_pipe = SDP.from_pretrained( args.model_version, use_auth_token=True, ) # Get Scheduler user_specified_scheduler = None if args.scheduler is not None: user_specified_scheduler = SCHEDULER_MAP[ args.scheduler].from_config(pytorch_pipe.scheduler.config) # Get Force Zeros Config if it exists force_zeros_for_empty_prompt: bool = False if 'xl' in args.model_version and 'force_zeros_for_empty_prompt' in pytorch_pipe.config: force_zeros_for_empty_prompt = pytorch_pipe.config['force_zeros_for_empty_prompt'] coreml_pipe = get_coreml_pipe( pytorch_pipe=pytorch_pipe, mlpackages_dir=args.i, model_version=args.model_version, compute_unit=args.compute_unit, scheduler_override=user_specified_scheduler, controlnet_models=args.controlnet, force_zeros_for_empty_prompt=force_zeros_for_empty_prompt, sources=args.model_sources, ) if args.controlnet: controlnet_cond = [] for i, _ in enumerate(args.controlnet): image_path = args.controlnet_inputs[i] image = prepare_controlnet_cond(image_path, coreml_pipe.height, coreml_pipe.width) controlnet_cond.append(image) else: controlnet_cond = None logger.info("Beginning image generation.") image = coreml_pipe( prompt=args.prompt, height=coreml_pipe.height, width=coreml_pipe.width, num_inference_steps=args.num_inference_steps, guidance_scale=args.guidance_scale, controlnet_cond=controlnet_cond, negative_prompt=args.negative_prompt, unet_batch_one=args.unet_batch_one, ) out_path = get_image_path(args) logger.info(f"Saving generated image to {out_path}") image["images"][0].save(out_path) if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument( "--prompt", required=True, help="The text prompt to be used for text-to-image generation.") parser.add_argument( "-i", required=True, help=("Path to input directory with the .mlpackage files generated by " "python_coreml_stable_diffusion.torch2coreml")) parser.add_argument("-o", required=True) parser.add_argument("--seed", "-s", default=93, type=int, help="Random seed to be able to reproduce results") parser.add_argument( "--model-version", default="CompVis/stable-diffusion-v1-4", help= ("The pre-trained model checkpoint and configuration to restore. " "For available versions: https://huggingface.co/models?search=stable-diffusion" )) parser.add_argument( "--compute-unit", choices=get_available_compute_units(), default="ALL", help=("The compute units to be used when executing Core ML models. " f"Options: {get_available_compute_units()}")) parser.add_argument( "--scheduler", choices=tuple(SCHEDULER_MAP.keys()), default=None, help=("The scheduler to use for running the reverse diffusion process. " "If not specified, the default scheduler from the diffusers pipeline is utilized")) parser.add_argument( "--num-inference-steps", default=50, type=int, help="The number of iterations the unet model will be executed throughout the reverse diffusion process") parser.add_argument( "--guidance-scale", default=7.5, type=float, help="Controls the influence of the text prompt on sampling process (0=random images)") parser.add_argument( "--controlnet", nargs="*", type=str, help=("Enables ControlNet and use control-unet instead of unet for additional inputs. " "For Multi-Controlnet, provide the model names separated by spaces.")) parser.add_argument( "--controlnet-inputs", nargs="*", type=str, help=("Image paths for ControlNet inputs. " "Please enter images corresponding to each controlnet provided at --controlnet option in same order.")) parser.add_argument( "--negative-prompt", default=None, help="The negative text prompt to be used for text-to-image generation.") parser.add_argument( "--unet-batch-one", action="store_true", help="Do not batch unet predictions for the prompt and negative prompt.") parser.add_argument('--model-sources', default=None, choices=['packages', 'compiled'], help='Force build from `packages` or `compiled`') args = parser.parse_args() main(args)