Merge pull request #217 from rupeshs/add-aura-sr-support

Add aura SR support
rupeshs · Jun 29, 2024 · 3ca4554 · 3ca4554
2 parents 30ea900 + ad224f0
commit 3ca4554
Show file tree

Hide file tree

Showing 10 changed files with 861 additions and 5 deletions.
diff --git a/Readme.md b/Readme.md
@@ -117,6 +117,7 @@ If we enable Tiny decoder(TAESD) we can save some memory(2GB approx) for example
 - 1 step fast inference support for SDXL and SD1.5
 - Experimental support for single file Safetensors SD 1.5 models(Civitai models), simply add local model path to configs/stable-diffusion-models.txt file.
 - Add REST API support
+- Add Aura SR (4x)/GigaGAN based upscaler support
 
 <a id="fast-inference-benchmarks"></a>
 

diff --git a/requirements.txt b/requirements.txt
@@ -10,7 +10,7 @@ onnxruntime==1.17.3
 pydantic==2.4.2
 typing-extensions==4.8.0
 pyyaml==6.0.1
-gradio==4.21.0
+gradio==4.23.0
 peft==0.6.1
 opencv-python==4.8.1.78
 omegaconf==2.3.0

diff --git a/src/backend/lcm_text_to_image.py b/src/backend/lcm_text_to_image.py
@@ -83,6 +83,9 @@ def _add_freeu(self):
                     b2=1.2,
                 )
 
+    def _enable_vae_tiling(self):
+        self.pipeline.vae.enable_tiling()
+
     def _update_lcm_scheduler_params(self):
         if isinstance(self.pipeline.scheduler, LCMScheduler):
             self.pipeline.scheduler = LCMScheduler.from_config(

diff --git a/src/backend/models/upscale.py b/src/backend/models/upscale.py
@@ -6,3 +6,4 @@ class UpscaleMode(str, Enum):
 
     normal = "normal"
     sd_upscale = "sd_upscale"
+    aura_sr = "aura_sr"