Update app.py

app.py

@@ -6,12 +6,17 @@
 # copies of the Software, and to permit persons to whom the Software is
 import spaces
 import os
+
+os.putenv('PYTORCH_NVML_BASED_CUDA_CHECK','1')
 os.putenv('TORCH_LINALG_PREFER_CUSOLVER','1')
-os.putenv('PYTORCH_CUDA_ALLOC_CONF','max_split_size_mb:128')
+alloc_conf_parts = [
+    'expandable_segments:True',
+    'pinned_use_background_threads:True'  # Specific to pinned memory.
+]
+os.environ['PYTORCH_CUDA_ALLOC_CONF'] = ','.join(alloc_conf_parts)
 os.environ["SAFETENSORS_FAST_GPU"] = "1"
 os.putenv('HF_HUB_ENABLE_HF_TRANSFER','1')
 
-
 import random
 import uuid
 import gradio as gr
@@ -285,7 +290,8 @@ def generate_30(
     #gc.collect()
     torch.set_float32_matmul_precision("medium")
     with torch.no_grad():
-        upscale = upscaler(rv_image, tiling=True, tile_width=256, tile_height=256)
+        upscalea = upscaler(rv_image, tiling=True, tile_width=256, tile_height=256)
+        upscale = upscaler(upscalea, tiling=True, tile_width=256, tile_height=256)
     downscale1 = upscale.resize((upscale.width // 4, upscale.height // 4), Image.LANCZOS)
     downscale_path = f"rv50E_upscale_{timestamp}.png"
     downscale1.save(downscale_path,optimize=False,compress_level=0)
@@ -340,7 +346,8 @@ def generate_60(
     image_paths = save_image(rv_image)
     torch.set_float32_matmul_precision("medium")
     with torch.no_grad():
-        upscale = upscaler(rv_image, tiling=True, tile_width=256, tile_height=256)
+        upscalea = upscaler(rv_image, tiling=True, tile_width=256, tile_height=256)
+        upscale = upscaler(upscalea, tiling=True, tile_width=256, tile_height=256)
     downscale1 = upscale.resize((upscale.width // 4, upscale.height // 4), Image.LANCZOS)
     downscale_path = f"rv50E_upscale_{timestamp}.png"
     downscale1.save(downscale_path,optimize=False,compress_level=0)
@@ -395,7 +402,8 @@ def generate_90(
     image_paths = save_image(rv_image)
     torch.set_float32_matmul_precision("medium")
     with torch.no_grad():
-        upscale = upscaler(rv_image, tiling=True, tile_width=256, tile_height=256)
+        upscalea = upscaler(rv_image, tiling=True, tile_width=256, tile_height=256)
+        upscale = upscaler(upscalea, tiling=True, tile_width=256, tile_height=256)
     downscale1 = upscale.resize((upscale.width // 4, upscale.height // 4), Image.LANCZOS)
     downscale_path = f"rv50E_upscale_{timestamp}.png"
     downscale1.save(downscale_path,optimize=False,compress_level=0)