Add RealESRGAN upscalling option

requirements.txt

@@ -3,4 +3,6 @@ diffusers
 omegaconf
 transformers
 peft
-setuptools
+setuptools
+torchsde
+git+https://github.com/sberbank-ai/Real-ESRGAN.git

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={},
-)
+)

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 = []

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,19 +20,20 @@ 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.
 
     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

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")

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()