Update app.py

app.py

@@ -27,6 +27,8 @@ import datetime
 from gradio import themes
 from hidiffusion import apply_hidiffusion, remove_hidiffusion
 
+import gc
+
 torch.backends.cuda.matmul.allow_tf32 = False
 torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False
 torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False
@@ -161,7 +163,7 @@ def load_and_prepare_model(model_id):
     print(f'watermark: {pipe.watermark}')
     print(f'image_processor: {pipe.image_processor}')
     print(f'feature_extractor: {pipe.feature_extractor}')
-    #print(f'VAE: {pipe.vae}')
+    print(f'init noise scale: {pipe.scheduler.init_noise_sigma}')
     #print(f'UNET: {pipe.unet}')
     pipe.watermark=None
     pipe.safety_checker=None  
@@ -221,6 +223,8 @@ def generate_30(
     num_images: int = 1,  
     progress=gr.Progress(track_tqdm=True)  # Add progress as a keyword argument
 ):
+    torch.cuda.empty_cache()
+    gc.collect()
     global models
     pipe = models[model_choice]
     seed = int(randomize_seed_fn(seed, randomize_seed))
@@ -288,6 +292,8 @@ def generate_60(
     num_images: int = 1,
     progress=gr.Progress(track_tqdm=True)  # Add progress as a keyword argument
 ):
+    torch.cuda.empty_cache()
+    gc.collect()
     global models
     pipe = models[model_choice]
     seed = int(randomize_seed_fn(seed, randomize_seed))
@@ -355,6 +361,8 @@ def generate_90(
     num_images: int = 1,  
     progress=gr.Progress(track_tqdm=True)  # Add progress as a keyword argument
 ):
+    torch.cuda.empty_cache()
+    gc.collect()
     global models
     pipe = models[model_choice]
     seed = int(randomize_seed_fn(seed, randomize_seed))