From c5544764dd5e4ee5235b7d5a199f2f77c3cabacc Mon Sep 17 00:00:00 2001 From: Leonardo Severo Date: Mon, 1 Jul 2024 10:11:11 -0300 Subject: [PATCH 1/3] Adding upscaling and changing pipeline scheduler --- requirements.txt | 4 +++- setup.py | 7 +++++-- sevsd/generate_image.py | 33 +++++++++++++++++++++++++++------ sevsd/setup_pipeline.py | 8 ++++++-- 4 files changed, 41 insertions(+), 11 deletions(-) diff --git a/requirements.txt b/requirements.txt index e2f1fc1..ac04cd2 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,4 +3,6 @@ diffusers omegaconf transformers peft -setuptools \ No newline at end of file +setuptools +torchsde +git+https://github.com/sberbank-ai/Real-ESRGAN.git \ No newline at end of file diff --git a/setup.py b/setup.py index e92a0e2..0f9e01b 100644 --- a/setup.py +++ b/setup.py @@ -22,9 +22,12 @@ "transformers", "omegaconf", "peft>=0.11.0", - "setuptools" + "setuptools", + "RealESRGAN", + "torchsde", + "realesrgan @ git+https://github.com/sberbank-ai/Real-ESRGAN.git", ], include_package_data=True, keywords='image generation, stable diffusion, AI', entry_points={}, -) +) \ No newline at end of file diff --git a/sevsd/generate_image.py b/sevsd/generate_image.py index 275e145..a068ef1 100644 --- a/sevsd/generate_image.py +++ b/sevsd/generate_image.py @@ -1,8 +1,14 @@ +import numpy as np import torch +from PIL import Image +from RealESRGAN import RealESRGAN +from sevsd.setup_device import setup_device -def generate_image(job, pipeline, executor, parallel_exec=True, **kwargs): +device = setup_device() + +def generate_image(job, pipeline, executor, upscaling_factor=4, parallel_exec=True, **kwargs): r""" - Generates images based on textual prompts using the provided Stable Diffusion pipeline. + Generates high-resolution images based on textual prompts using the provided Stable Diffusion pipeline and an upscaling model. This function handles the image generation process using specified parameters and execution configurations. It supports both parallel and sequential execution modes for image generation tasks. @@ -16,11 +22,12 @@ def generate_image(job, pipeline, executor, parallel_exec=True, **kwargs): - 'inference_steps' (int): Number of inference steps for the pipeline. - 'num_of_exec' (int): Number of images to generate. - 'cfg_scale' (float): The guidance scale for controlling image generation. + upscaling_factor (int, optional): The factor by which to upscale the generated images. Defaults to 4. parallel_exec (bool, optional): If True, generates 'num_of_exec' images in parallel. Otherwise, generates images sequentially. Defaults to True. **kwargs: Additional keyword arguments passed to the pipeline function. Returns: - list or None: A list of generated images if successful, or None if a RuntimeError occurs during image generation. + list or None: A list of generated high-resolution images if successful, or None if a RuntimeError occurs during image generation. Raises: RuntimeError: If an error occurs during the image generation process. @@ -28,7 +35,7 @@ def generate_image(job, pipeline, executor, parallel_exec=True, **kwargs): Example: >>> job = {"prompt": "A beautiful landscape", "negative_prompt": "dark, blurry images"} >>> executor = {"inference_steps": 50, "num_of_exec": 3, "cfg_scale": 7.5} - >>> images = generate_image(job, pipeline, executor, parallel_exec=True) + >>> images = generate_image(job, pipeline, executor, upscaling_factor=4, parallel_exec=True) >>> len(images) # Should be equal to the number specified in 'num_of_exec' """ prompt = job.get("prompt") @@ -48,14 +55,28 @@ def execute_pipeline(num_images): negative_embeddings=getattr(pipeline, 'negative_embeddings', None), **kwargs ) + + def upscale_image(image, factor): + model = RealESRGAN(device, scale=factor) + model.load_weights('weights/RealESRGAN_x4.pth', download=True) + + pil_image = Image.fromarray(image) + upscaled_image = model.predict(pil_image) + + return np.array(upscaled_image) + try: with torch.no_grad(): if parallel_exec: output = execute_pipeline(num_images) - return output["images"] + base_images = output["images"] else: output = execute_pipeline(1) - return output["images"] + base_images = output["images"] + + high_res_images = [upscale_image(img, upscaling_factor) for img in base_images] + + return high_res_images except RuntimeError as e: print(f"Runtime error: {e}") return None \ No newline at end of file diff --git a/sevsd/setup_pipeline.py b/sevsd/setup_pipeline.py index d85123f..27fbd87 100644 --- a/sevsd/setup_pipeline.py +++ b/sevsd/setup_pipeline.py @@ -1,5 +1,5 @@ from sevsd.setup_device import setup_device -from diffusers import StableDiffusionPipeline, EulerAncestralDiscreteScheduler +from diffusers import StableDiffusionPipeline, DPMSolverSDEScheduler from transformers import AutoFeatureExtractor import os @@ -51,7 +51,11 @@ def setup_pipeline(pretrained_model_link_or_path, loras, positive_embeddings=Non **default_kwargs ) - pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(pipeline.scheduler.config) + pipeline.scheduler = DPMSolverSDEScheduler.from_config( + pipeline.scheduler.config, + algorithm_type="karras", + use_advanced_sde=True + ) if loras: pipeline.unfuse_lora() From d46cfc29535bdd6d78ca54babbe044c2bce0f233 Mon Sep 17 00:00:00 2001 From: Leonardo Severo Date: Mon, 1 Jul 2024 10:41:57 -0300 Subject: [PATCH 2/3] Removing imports --- sevsd/do_work.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/sevsd/do_work.py b/sevsd/do_work.py index d9ce012..914dbd5 100644 --- a/sevsd/do_work.py +++ b/sevsd/do_work.py @@ -1,8 +1,6 @@ -import torch -import os from sevsd.setup_pipeline import setup_pipeline from sevsd.process_task import process_task -from sevsd.load_embeddings import load_embeddings_from_folder # Assuming load_embeddings_from_folder is saved here +from sevsd.load_embeddings import load_embeddings_from_folder def load_all_embeddings(folders): embeddings = [] From 3ca84b2b9ba8008c41f60cb4e2a0cd9aea2320c0 Mon Sep 17 00:00:00 2001 From: Leonardo Severo Date: Mon, 1 Jul 2024 19:04:14 -0300 Subject: [PATCH 3/3] Add upscaling support --- sevsd/generate_image.py | 32 ++++++++++++++++++++------------ sevsd/process_task.py | 2 +- 2 files changed, 21 insertions(+), 13 deletions(-) diff --git a/sevsd/generate_image.py b/sevsd/generate_image.py index a068ef1..3b47721 100644 --- a/sevsd/generate_image.py +++ b/sevsd/generate_image.py @@ -1,12 +1,10 @@ -import numpy as np import torch -from PIL import Image from RealESRGAN import RealESRGAN from sevsd.setup_device import setup_device device = setup_device() -def generate_image(job, pipeline, executor, upscaling_factor=4, parallel_exec=True, **kwargs): +def generate_image(job, pipeline, executor, upscaling_factor=2, parallel_exec=True, **kwargs): r""" Generates high-resolution images based on textual prompts using the provided Stable Diffusion pipeline and an upscaling model. @@ -22,7 +20,7 @@ def generate_image(job, pipeline, executor, upscaling_factor=4, parallel_exec=Tr - 'inference_steps' (int): Number of inference steps for the pipeline. - 'num_of_exec' (int): Number of images to generate. - 'cfg_scale' (float): The guidance scale for controlling image generation. - upscaling_factor (int, optional): The factor by which to upscale the generated images. Defaults to 4. + upscaling_factor (int, optional): The factor by which to upscale the generated images. Defaults to 2. parallel_exec (bool, optional): If True, generates 'num_of_exec' images in parallel. Otherwise, generates images sequentially. Defaults to True. **kwargs: Additional keyword arguments passed to the pipeline function. @@ -57,13 +55,19 @@ def execute_pipeline(num_images): ) def upscale_image(image, factor): - model = RealESRGAN(device, scale=factor) - model.load_weights('weights/RealESRGAN_x4.pth', download=True) + if factor not in [2, 4, 8]: + raise ValueError("Upscaling factor must be one of the following: 2, 4, 8") - pil_image = Image.fromarray(image) - upscaled_image = model.predict(pil_image) + model = RealESRGAN(device, scale=factor) + try: + model.load_weights('weights/RealESRGAN_x{}.pth'.format(factor), download=True) + except Exception as e: + print(f"Failed to load RealESRGAN weights: {e}") + return None - return np.array(upscaled_image) + upscaled_image = model.predict(image) + + return upscaled_image try: with torch.no_grad(): @@ -73,9 +77,13 @@ def upscale_image(image, factor): else: output = execute_pipeline(1) base_images = output["images"] - - high_res_images = [upscale_image(img, upscaling_factor) for img in base_images] - + high_res_images = [] + for img in base_images: + upscaled_img = upscale_image(img, upscaling_factor) + if upscaled_img is not None: + high_res_images.append(upscaled_img) + else: + high_res_images.append(img) return high_res_images except RuntimeError as e: print(f"Runtime error: {e}") diff --git a/sevsd/process_task.py b/sevsd/process_task.py index ab04691..9417aef 100644 --- a/sevsd/process_task.py +++ b/sevsd/process_task.py @@ -39,7 +39,7 @@ def process_task(job, pipeline, executor, path, parallel_exec=True): """ def call_generate_image(): - images = generate_image(job, pipeline, executor, parallel_exec) + images = generate_image(job, pipeline, executor, 2, parallel_exec) if images is not None: for image in images: timestamp = datetime.now().strftime("%Y%m%d_%H%M%S%f")