Update app.py

app.py

@@ -5,7 +5,6 @@
 # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 # copies of the Software, and to permit persons to whom the Software is
 import spaces
-
 import os
 import random
 import uuid
@@ -23,10 +22,8 @@ import gc
 import time
 import datetime
 #from diffusers.schedulers import AysSchedules
-
 from gradio import themes
 from hidiffusion import apply_hidiffusion, remove_hidiffusion
-
 import gc
 
 torch.backends.cuda.matmul.allow_tf32 = False
@@ -37,7 +34,6 @@ torch.backends.cudnn.deterministic = False
 #torch.backends.cudnn.benchmark = False
 torch.backends.cuda.preferred_blas_library="cublas"
 # torch.backends.cuda.preferred_linalg_library="cusolver"
-
 torch.set_float32_matmul_precision("highest")
 
 FTP_HOST = "1ink.us"
@@ -50,7 +46,6 @@ DESCRIPTIONXX = """
 """
 
 examples = [
-
     "Many apples splashed with drops of water within a fancy bowl 4k, hdr  --v 6.0 --style raw",
     "A profile photo of a dog, brown background, shot on Leica M6 --ar 128:85 --v 6.0 --style raw",
 ]
@@ -60,12 +55,8 @@ MODEL_OPTIONS = {
 }
 
 MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "4096"))
-USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE", "0") == "1"
-ENABLE_CPU_OFFLOAD = 0
 BATCH_SIZE = int(os.getenv("BATCH_SIZE", "1"))
 
-device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-
 style_list = [
     {
         "name": "3840 x 2160",
@@ -95,15 +86,7 @@ STYLE_NAMES = list(styles.keys())
 HF_TOKEN = os.getenv("HF_TOKEN")
 
 #sampling_schedule = AysSchedules["StableDiffusionXLTimesteps"]
-
-def apply_style(style_name: str, positive: str, negative: str = "") -> Tuple[str, str]:
-    if style_name in styles:
-        p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
-    else:
-        p, n = styles[DEFAULT_STYLE_NAME]
-    if not negative:
-        negative = ""
-    return p.replace("{prompt}", positive), n + negative
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
     
 def load_and_prepare_model(model_id):
     model_dtypes = {"ford442/RealVisXL_V5.0_BF16": torch.bfloat16,}
@@ -200,8 +183,15 @@ def upload_to_ftp(filename):
         print(f"Uploaded {filename} to FTP server")
     except Exception as e:
         print(f"FTP upload error: {e}")
-        
-def save_image(img):
+
+def apply_style(style_name: str, positive: str, negative: str = "") -> Tuple[str, str]:
+    if style_name in styles:
+        p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
+    else:
+        p, n = styles[DEFAULT_STYLE_NAME]
+    if not negative:
+        negative = ""
+    return p.replace("{prompt}", positive), n + negativedef save_image(img):
     unique_name = str(uuid.uuid4()) + ".png"
     img.save(unique_name,optimize=False,compress_level=0)
     return unique_name
@@ -246,6 +236,7 @@ def generate_30(
     num_images: int = 1,  
     progress=gr.Progress(track_tqdm=True)  # Add progress as a keyword argument
 ):
+    torch.backends.cudnn.benchmark = False
     torch.cuda.empty_cache()
     gc.collect()
     global models
@@ -302,6 +293,7 @@ def generate_60(
     num_images: int = 1,
     progress=gr.Progress(track_tqdm=True)  # Add progress as a keyword argument
 ):
+    torch.backends.cudnn.benchmark = True
     torch.cuda.empty_cache()
     gc.collect()
     global models
@@ -358,6 +350,7 @@ def generate_90(
     num_images: int = 1,  
     progress=gr.Progress(track_tqdm=True)  # Add progress as a keyword argument
 ):
+    torch.backends.cudnn.benchmark = True
     torch.cuda.empty_cache()
     gc.collect()
     global models