diff --git a/models/model_text_to_image.py b/models/model_text_to_image.py index 12f6763..35690c3 100644 --- a/models/model_text_to_image.py +++ b/models/model_text_to_image.py @@ -64,10 +64,12 @@ def unload_model(self) -> bool: self.loaded = False return True - def generate_prompt(self, prompt: Optional[str], options: OptionsTextToImage): + def generate_prompt(self, prompt: Optional[str], + options: OptionsTextToImage): """ Generates the prompt with the given option - :param prompt: The optional prompt (if the prompt is empty, the options.prompt will be used) + :param prompt: The optional prompt + (if the prompt is empty, the options.prompt will be used) :param options: The options of text to image model :return: An object image resulting from the model """ @@ -89,7 +91,8 @@ def generate_prompt(self, prompt: Optional[str], options: OptionsTextToImage): prompt_embeds=options.prompt_embeds, negative_prompt_embeds=options.negative_prompt_embeds, pooled_prompt_embeds=options.pooled_prompt_embeds, - negative_pooled_prompt_embeds=options.negative_pooled_prompt_embeds, + negative_pooled_prompt_embeds=( + options.negative_pooled_prompt_embeds), ip_adapter_image=options.ip_adapter_image, output_type=options.output_type, return_dict=options.return_dict, @@ -99,9 +102,11 @@ def generate_prompt(self, prompt: Optional[str], options: OptionsTextToImage): crops_coords_top_left=options.crops_coords_top_left, target_size=options.target_size, negative_original_size=options.negative_original_size, - negative_crops_coords_top_left=options.negative_crops_coords_top_left, + negative_crops_coords_top_left=( + options.negative_crops_coords_top_left), negative_target_size=options.negative_target_size, clip_skip=options.clip_skip, callback_on_step_end=options.callback_on_step_end, - callback_on_step_end_tensor_inputs=options.callback_on_step_end_tensor_inputs + callback_on_step_end_tensor_inputs=( + options.callback_on_step_end_tensor_inputs) ).images[0] diff --git a/models/models_management.py b/models/models_management.py index 3e3acea..4d8bc60 100644 --- a/models/models_management.py +++ b/models/models_management.py @@ -6,7 +6,8 @@ class ModelsManagement: """ The ModelsManagement class controls all instantiated models. - It is with this class that you can deploy a model on a device and generate a prompt. + It is with this class that you can deploy a model on a device and + generate a prompt. """ def __init__(self): @@ -39,7 +40,8 @@ def load_model(self, model_name: str) -> bool: :return: True if the model is successfully loaded """ if self.loaded_model: - print("Unload the currently loaded model before loading a new one.") + print( + "Unload the currently loaded model before loading a new one.") return False if model_name not in self.loaded_models_cache: @@ -47,7 +49,8 @@ def load_model(self, model_name: str) -> bool: return False self.loaded_model = self.loaded_models_cache[model_name] - if not self.loaded_model.load_model(option=self.options_models[model_name]): + if not self.loaded_model.load_model(option=( + self.options_models[model_name])): print("Something went wrong while unloading the model.") self.loaded_model = None return False @@ -88,14 +91,20 @@ def set_model_options(self, model_name: str, options: Options): def generate_prompt(self, prompt: Optional[str] = None): """ Generates the prompt for the loaded model with his stored options - :param prompt: The prompt to generate (if the prompt is empty, the options.prompt will be used) + :param prompt: The prompt to generate (if the prompt is empty, the + options.prompt will be used) :return: The object of type link with the model category """ if not self.loaded_model: print("No model loaded. Load a model before generating prompts.") return - return self.loaded_model.generate_prompt(prompt, self.options_models[self.loaded_model.model_name]) + return ( + self.loaded_model.generate_prompt(prompt, + self.options_models[ + self.loaded_model.model_name] + ) + ) def print_models(self): """ @@ -103,5 +112,6 @@ def print_models(self): """ print("Models in cache:") for model_name, model_instance in self.loaded_models_cache.items(): - selected_indicator = "(selected)" if model_instance == self.loaded_model else "" + selected_indicator = ( + "(selected)" if model_instance == self.loaded_model else "") print(f"- {model_name} {selected_indicator}") diff --git a/options/options_text_to_image.py b/options/options_text_to_image.py index c0ff2d7..c08857e 100644 --- a/options/options_text_to_image.py +++ b/options/options_text_to_image.py @@ -56,7 +56,8 @@ def __init__( negative_prompt_2: Optional[Union[str, List[str]]] = None, num_images_per_prompt: Optional[int] = None, eta: Optional[float] = None, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + generator: Optional[ + Union[torch.Generator, List[torch.Generator]]] = None, latents: Optional[torch.FloatTensor] = None, prompt_embeds: Optional[torch.FloatTensor] = None, negative_prompt_embeds: Optional[torch.FloatTensor] = None, @@ -74,7 +75,8 @@ def __init__( negative_crops_coords_top_left: Tuple[int, int] = None, negative_target_size: Optional[Tuple[int, int]] = None, clip_skip: Optional[int] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, + callback_on_step_end: Optional[ + Callable[[int, int, Dict], None]] = None, callback_on_step_end_tensor_inputs: List[str] = None ): """ @@ -82,101 +84,138 @@ def __init__( :param device: The device to use generate prompt :param prompt: The prompt to give to the model :param image_width: The width of the resulting image (1024 by default) - :param image_height: The height of the resulting image (1024 by default) + :param image_height: The height of the resulting image + (1024 by default) :param num_inference_steps (`int`, *optional*, defaults to 50): - The number of denoising steps. More denoising steps usually lead to a higher quality image at the - expense of slower inference. + The number of denoising steps. More denoising steps usually lead to + a higher quality image at the expense of slower inference. :param timesteps (`List[int]`, *optional*): - Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument - in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is - passed will be used. Must be in descending order. + Custom timesteps to use for the denoising process with schedulers + which support a `timesteps` argument in their `set_timesteps` + method. If not defined, the default behavior when + `num_inference_steps` is passed will be used. + Must be in descending order. :param denoising_end (`float`, *optional*): - When specified, determines the fraction (between 0.0 and 1.0) of the total denoising process to be - completed before it is intentionally prematurely terminated. As a result, the returned sample will - still retain a substantial amount of noise as determined by the discrete timesteps selected by the - scheduler. The denoising_end parameter should ideally be utilized when this pipeline forms a part of a + When specified, determines the fraction (between 0.0 and 1.0) of + the total denoising process to be completed before it is + intentionally prematurely terminated. + As a result, the returned sample will still retain a substantial + amount of noise as determined by the discrete timesteps selected by + the scheduler. The denoising_end parameter should ideally be + utilized when this pipeline forms a part of a "Mixture of Denoisers" multi-pipeline setup. :param guidance_scale (`float`, *optional*, defaults to 5.0): - Guidance scale as defined in Classifier-Free Diffusion Guidance. `guidance_scale` is defined as `w` - of equation 2. of the Imagen Paper. Guidance scale is enabled by setting `guidance_scale > 1`. - Higher guidance scale encourages generating images closely linked to the text `prompt`, usually at the - expense of lower image quality. + Guidance scale as defined in Classifier-Free Diffusion Guidance. + `guidance_scale` is defined as `w` of equation 2. of the Imagen + Paper. Guidance scale is enabled by setting `guidance_scale > 1`. + Higher guidance scale encourages generating images closely linked + to the text `prompt`, usually at the expense of lower image quality :param negative_prompt (`str` or `List[str]`, *optional*): - The prompt or prompts not to guide the image generation. If not defined, one has to pass - `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is - less than `1`). + The prompt or prompts not to guide the image generation. If not + defined, one has to pass `negative_prompt_embeds` instead. + Ignored when not using guidance (i.e., ignored if `guidance_scale` + is less than `1`). :param negative_prompt_2 (`str` or `List[str]`, *optional*): - The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and - `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders. + The prompt or prompts not to guide the image generation to be sent + to `tokenizer_2` and `text_encoder_2`. If not defined, + `negative_prompt` is used in both text-encoders. :param num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. :param eta (`float`, *optional*, defaults to 0.0): - Corresponds to parameter eta (η) in the DDIM paper. Only applies to [`schedulers.DDIMScheduler`], - will be ignored for others. - :param generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + Corresponds to parameter eta (η) in the DDIM paper. Only applies to + [`schedulers.DDIMScheduler`], will be ignored for others. + :param generator (`torch.Generator` or `List[torch.Generator]`, + *optional*): One or a list of torch generators to make generation deterministic. :param latents (`torch.FloatTensor`, *optional*): - Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image - generation. Can be used to tweak the same generation with different prompts. If not provided, a latents - tensor will be generated by sampling using the supplied random generator. + Pre-generated noisy latents, sampled from a Gaussian distribution, + to be used as inputs for image generation. Can be used to tweak the + same generation with different prompts. If not provided, a latents + tensor will be generated by sampling using the supplied random + generator. :param prompt_embeds (`torch.FloatTensor`, *optional*): - Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not - provided, text embeddings will be generated from `prompt` input argument. + Pre-generated text embeddings. Can be used to easily tweak text + inputs, *e.g.* prompt weighting. If not provided, text embeddings + will be generated from `prompt` input argument. :param negative_prompt_embeds (`torch.FloatTensor`, *optional*): - Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt - weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input - argument. + Pre-generated negative text embeddings. Can be used to easily tweak + text inputs, *e.g.* prompt weighting. If not provided, + negative_prompt_embeds will be generated from `negative_prompt` + input argument. :param pooled_prompt_embeds (`torch.FloatTensor`, *optional*): - Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. - If not provided, pooled text embeddings will be generated from `prompt` input argument. + Pre-generated pooled text embeddings. Can be used to easily tweak + text inputs, *e.g.* prompt weighting. If not provided, pooled text + embeddings will be generated from `prompt` input argument. :param negative_pooled_prompt_embeds (`torch.FloatTensor`, *optional*): - Pre-generated negative pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt - weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt` - input argument. - :param ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. + Pre-generated negative pooled text embeddings. + Can be used to easily tweak text inputs, *e.g.* prompt weighting. + If not provided, pooled negative_prompt_embeds will be generated + from `negative_prompt` input argument. + :param ip_adapter_image: (`PipelineImageInput`, *optional*): + Optional image input to work with IP Adapters. :param output_type (`str`, *optional*, defaults to `"pil"`): - The output format of the generated image. Choose between PIL.Image.Image or np.array. + The output format of the generated image. Choose between + PIL.Image.Image or np.array. :param return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a StableDiffusionXLPipelineOutput instead of a plain tuple. + Whether or not to return a StableDiffusionXLPipelineOutput instead + of a plain tuple. :param cross_attention_kwargs (`dict`, *optional*): - A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under - `self.processor` in diffusers.models.attention_processor. + A kwargs dictionary that if specified is passed along to the + `AttentionProcessor` as defined under `self.processor` in + diffusers.models.attention_processor. :param guidance_rescale (`float`, *optional*, defaults to 0.0): - Guidance rescale factor proposed by Common Diffusion Noise Schedules and Sample Steps are Flawed. - `guidance_scale` is defined as `φ` in equation 16. of Common Diffusion Noise Schedules and Sample Steps - are Flawed. Guidance rescale factor should fix overexposure when using zero terminal SNR. - :param original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): - If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled. - `original_size` defaults to `(height, width)` if not specified. Part of SDXL's micro-conditioning as + Guidance rescale factor proposed by Common Diffusion Noise + Schedules and Sample Steps are Flawed. `guidance_scale` is defined + as `φ` in equation 16. of Common Diffusion Noise Schedules and + Sample Steps are Flawed. Guidance rescale factor should fix + overexposure when using zero terminal SNR. + :param original_size + (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + If `original_size` is not the same as `target_size` the image will + appear to be down- or upsampled. `original_size` defaults to + `(height, width)` if not specified. Part of SDXL's + micro-conditioning as explained in section 2.2 of the documentation + :param crops_coords_top_left + (`Tuple[int]`, *optional*, defaults to (0, 0)): + `crops_coords_top_left` can be used to generate an image that + appears to be "cropped" from the position `crops_coords_top_left` + downwards. Favorable, well-centered images are usually achieved by + setting `crops_coords_top_left` to (0, 0). Part of SDXL's + micro-conditioning as explained in section 2.2 of the documentation + :param target_size + (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + For most cases, `target_size` should be set to the desired height + and width of the generated image. If not specified it will default + to `(height, width)`. Part of SDXL's micro-conditioning as explained in section 2.2 of the documentation. - :param crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)): - `crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position - `crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting - `crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of - the documentation. - :param target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): - For most cases, `target_size` should be set to the desired height and width of the generated image. If - not specified it will default to `(height, width)`. Part of SDXL's micro-conditioning as explained in + :param negative_original_size + (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + To negatively condition the generation process based on a specific + image resolution. Part of SDXL's micro-conditioning as explained in + section 2.2 of the documentation. + :param negative_crops_coords_top_left + (`Tuple[int]`, *optional*, defaults to (0, 0)): + To negatively condition the generation process based on specific + crop coordinates. Part of SDXL's micro-conditioning as explained in + section 2.2 of the documentation. + :param negative_target_size + (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + To negatively condition the generation process based on a target + image resolution. It should be as same as the `target_size` for + most cases. Part of SDXL's micro-conditioning as explained in section 2.2 of the documentation. - :param negative_original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): - To negatively condition the generation process based on a specific image resolution. Part of SDXL's - micro-conditioning as explained in section 2.2 of the documentation. - :param negative_crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)): - To negatively condition the generation process based on specific crop coordinates. Part of SDXL's - micro-conditioning as explained in section 2.2 of the documentation. - :param negative_target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): - To negatively condition the generation process based on a target image resolution. It should be as same - as the `target_size` for most cases. Part of SDXL's micro-conditioning as explained in section 2.2 of - the documentation. :param callback_on_step_end (`Callable`, *optional*): - A function that calls at the end of each denoising step during the inference. The function is called - with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, - callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by - `callback_on_step_end_tensor_inputs`. + A function that calls at the end of each denoising step during the + inference. The function is called with the following arguments: + `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: + int, callback_kwargs: Dict)`. `callback_kwargs` will include a list + of all tensors as specified by `callback_on_step_end_tensor_inputs` :param callback_on_step_end_tensor_inputs (`List`, *optional*): - The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list - will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the - `._callback_tensor_inputs` attribute of your pipeline class. + The list of tensor inputs for the `callback_on_step_end` function. + The tensors specified in the list will be passed as + `callback_kwargs` argument. You will only be able to include + variables listed in the `._callback_tensor_inputs` attribute of + your pipeline class. """ super().__init__(device) self.prompt = prompt @@ -215,9 +254,11 @@ def __init__( self.target_size = target_size self.negative_original_size = negative_original_size if negative_crops_coords_top_left: - self.negative_crops_coords_top_left = negative_crops_coords_top_left + self.negative_crops_coords_top_left = ( + negative_crops_coords_top_left) self.negative_target_size = negative_target_size self.clip_skip = clip_skip self.callback_on_step_end = callback_on_step_end if callback_on_step_end_tensor_inputs: - self.callback_on_step_end_tensor_inputs = callback_on_step_end_tensor_inputs + self.callback_on_step_end_tensor_inputs = ( + callback_on_step_end_tensor_inputs) diff --git a/tests/test.py b/tests/test.py index e74c6c8..f1448cc 100644 --- a/tests/test.py +++ b/tests/test.py @@ -1,8 +1,10 @@ import unittest + def add(a, b): return a + b + class TestAddFunction(unittest.TestCase): def test_add_positive_numbers(self): @@ -17,5 +19,6 @@ def test_add_mixed_numbers(self): result = add(10, -7) self.assertEqual(result, 3) + if __name__ == '__main__': - unittest.main() \ No newline at end of file + unittest.main()