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/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 = [] diff --git a/sevsd/generate_image.py b/sevsd/generate_image.py index 275e145..3b47721 100644 --- a/sevsd/generate_image.py +++ b/sevsd/generate_image.py @@ -1,8 +1,12 @@ import torch +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=2, 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 +20,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 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. 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 +33,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 +53,38 @@ def execute_pipeline(num_images): negative_embeddings=getattr(pipeline, 'negative_embeddings', None), **kwargs ) + + def upscale_image(image, factor): + if factor not in [2, 4, 8]: + raise ValueError("Upscaling factor must be one of the following: 2, 4, 8") + + 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 + + upscaled_image = model.predict(image) + + return 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 = [] + 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}") return None \ No newline at end of file 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") 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()