Update app.py

app.py

@@ -28,6 +28,7 @@ from gradio import themes
 from hidiffusion import apply_hidiffusion, remove_hidiffusion
 
 import gc
+from image_gen_aux import UpscaleWithModel
 
 torch.backends.cuda.matmul.allow_tf32 = False
 torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False
@@ -95,6 +96,7 @@ STYLE_NAMES = list(styles.keys())
 HF_TOKEN = os.getenv("HF_TOKEN")
 
 #sampling_schedule = AysSchedules["StableDiffusionXLTimesteps"]
+upscaler = UpscaleWithModel.from_pretrained("Kim2091/ClearRealityV1").to(torch.device("cuda:0"))
 
 def apply_style(style_name: str, positive: str, negative: str = "") -> Tuple[str, str]:
     if style_name in styles:
@@ -308,10 +310,17 @@ def generate_30(
         images.extend(pipe(**batch_options).images)
     sd_image_path = f"rv50_A_{timestamp}.png"
     images[0].save(sd_image_path,optimize=False,compress_level=0)
-    upload_to_ftp(sd_image_path)    
+    upload_to_ftp(sd_image_path)
     image_paths = [save_image(img) for img in images]
     torch.cuda.empty_cache()
     gc.collect()
+    torch.set_float32_matmul_precision("medium")
+    with torch.no_grad():
+        upscale = upscaler(images[0], tiling=True, tile_width=256, tile_height=256)
+    downscale1 = upscale.resize((upscale.width // 4, upscale.height // 4), Image.LANCZOS)
+        downscale_path = f"rv50_upscale_{timestamp}.png"
+    downscale1.save(downscale_path,optimize=False,compress_level=0)
+    upload_to_ftp(downscale_path)    
     return image_paths, seed
 
 @spaces.GPU(duration=60)
@@ -363,12 +372,19 @@ def generate_60(
         if "negative_prompt" in batch_options:
             batch_options["negative_prompt"] = options["negative_prompt"][i:i+BATCH_SIZE]
         images.extend(pipe(**batch_options).images)
-    sd_image_path = f"rv50_A_{seed}.png"
+    sd_image_path = f"rv50_A_{timestamp}.png"
     images[0].save(sd_image_path,optimize=False,compress_level=0)
-    upload_to_ftp(sd_image_path)    
+    upload_to_ftp(sd_image_path)
     image_paths = [save_image(img) for img in images]
     torch.cuda.empty_cache()
     gc.collect()
+    torch.set_float32_matmul_precision("medium")
+    with torch.no_grad():
+        upscale = upscaler(images[0], tiling=True, tile_width=256, tile_height=256)
+    downscale1 = upscale.resize((upscale.width // 4, upscale.height // 4), Image.LANCZOS)
+        downscale_path = f"rv50_upscale_{timestamp}.png"
+    downscale1.save(downscale_path,optimize=False,compress_level=0)
+    upload_to_ftp(downscale_path)    
     return image_paths, seed
 
 @spaces.GPU(duration=90)
@@ -422,10 +438,17 @@ def generate_90(
         images.extend(pipe(**batch_options).images)
     sd_image_path = f"rv50_A_{seed}.png"
     images[0].save(sd_image_path,optimize=False,compress_level=0)
-    upload_to_ftp(sd_image_path)    
+    upload_to_ftp(sd_image_path)
     image_paths = [save_image(img) for img in images]
     torch.cuda.empty_cache()
     gc.collect()
+    torch.set_float32_matmul_precision("medium")
+    with torch.no_grad():
+        upscale = upscaler(images[0], tiling=True, tile_width=256, tile_height=256)
+    downscale1 = upscale.resize((upscale.width // 4, upscale.height // 4), Image.LANCZOS)
+        downscale_path = f"rv50_upscale_{timestamp}.png"
+    downscale1.save(downscale_path,optimize=False,compress_level=0)
+    upload_to_ftp(downscale_path)    
     return image_paths, seed
 
 def load_predefined_images1():