Update app.py

app.py

@@ -20,6 +20,7 @@ import paramiko
 import datetime
 #from diffusers import DPMSolverSDEScheduler
 from diffusers.models.attention_processor import AttnProcessor2_0
+import gc
 
 torch.backends.cuda.matmul.allow_tf32 = False
 torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False
@@ -85,10 +86,10 @@ device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
     
 def load_and_prepare_model():
     #vaeRV = AutoencoderKL.from_pretrained("SG161222/RealVisXL_V5.0", subfolder='vae', safety_checker=None, use_safetensors=True, token=True)
-    vaeXL = AutoencoderKL.from_pretrained("stabilityai/sdxl-vae", safety_checker=None, use_safetensors=False) #.to(device).to(torch.bfloat16) #.to(device=device, dtype=torch.bfloat16)
+    vaeXL = AutoencoderKL.from_pretrained("stabilityai/sdxl-vae", safety_checker=None, use_safetensors=False).to(device) # .to(torch.bfloat16) #.to(device=device, dtype=torch.bfloat16)
     #sched = EulerAncestralDiscreteScheduler.from_pretrained('ford442/RealVisXL_V5.0_BF16', subfolder='scheduler',beta_schedule="scaled_linear", beta_start=0.00085, beta_end=0.012, steps_offset=1,use_karras_sigmas=True)
     #sched = DPMSolverSDEScheduler.from_pretrained('ford442/RealVisXL_V5.0_BF16', subfolder='scheduler')
-    sched = EulerAncestralDiscreteScheduler.from_pretrained('ford442/RealVisXL_V5.0_BF16', subfolder='scheduler',beta_schedule="scaled_linear", token=True) #, beta_start=0.00085, beta_end=0.012, steps_offset=1,use_karras_sigmas=True, token=True)
+    #sched = EulerAncestralDiscreteScheduler.from_pretrained('ford442/RealVisXL_V5.0_BF16', subfolder='scheduler',beta_schedule="scaled_linear", token=True) #, beta_start=0.00085, beta_end=0.012, steps_offset=1,use_karras_sigmas=True, token=True)
     #sched = EulerAncestralDiscreteScheduler.from_pretrained('ford442/RealVisXL_V5.0_BF16', subfolder='scheduler',beta_schedule="scaled_linear")
     pipe = StableDiffusionXLPipeline.from_pretrained(
         'ford442/RealVisXL_V5.0_BF16',
@@ -97,7 +98,6 @@ def load_and_prepare_model():
        # low_cpu_mem_usage = False,
         add_watermarker=False,
     )
-    pipe.vae = vaeXL #.to(torch.bfloat16)
     pipe.scheduler = sched
     #pipe.vae.do_resize=False
     #pipe.vae.vae_scale_factor=8
@@ -110,15 +110,8 @@ def load_and_prepare_model():
     #pipe.unet.to(memory_format=torch.channels_last)
     #pipe.enable_vae_tiling()
     pipe.to(device=device, dtype=torch.bfloat16)
-    '''
-    pipe.unet.set_attn_processor(AttnProcessor2_0())
-    for name, param in pipe.unet.named_parameters():
-        if 'qkv' in name or 'norm' in name or 'proj' in name:
-            param.requires_grad_()
-        else:
-            param.requires_grad_(False)
-    '''        
-    pipe.unet.requires_grad_(False)
+ 
+    pipe.vae = vaeXL #.to(torch.bfloat16)
     return pipe
     
 pipe = load_and_prepare_model()
@@ -171,6 +164,11 @@ def uploadNote(prompt,num_inference_steps,guidance_scale,timestamp):
         f.write(f"To cuda and bfloat \n")
     upload_to_ftp(filename) 
     
+def flush():
+  gc.collect()
+  torch.cuda.empty_cache()
+  torch.cuda.reset_peak_memory_stats()
+
 @spaces.GPU(duration=30)
 def generate_30(
     prompt: str,
@@ -186,10 +184,24 @@ def generate_30(
 ):
     seed = random.randint(0, MAX_SEED)
     generator = torch.Generator(device='cuda').manual_seed(seed)
+
+    #  get prompt embeds first
+    (
+    prompt_embeds,
+    negative_prompt_embeds,
+    pooled_prompt_embeds,
+    negative_pooled_prompt_embeds
+    ) = pipe.encode_prompt(prompt)
+    del pipe.text_encoder, pipe.text_encoder_2
+    flush()
     options = {
-        "prompt": [prompt],
-        "negative_prompt": [negative_prompt],
-        "negative_prompt_2": [neg_prompt_2],
+     #   "prompt": [prompt],
+      #  "negative_prompt": [negative_prompt],
+     #   "negative_prompt_2": [neg_prompt_2],
+		"prompt_embeds": = [prompt_embeds],
+		"negative_prompt_embeds": = [negative_prompt_embeds],
+		"pooled_prompt_embeds": = [pooled_prompt_embeds],
+		"negative_pooled_prompt_embeds": = [negative_pooled_prompt_embeds],
         "width": width,
         "height": height,
         "guidance_scale": guidance_scale,
@@ -226,10 +238,24 @@ def generate_60(
 ):
     seed = random.randint(0, MAX_SEED)
     generator = torch.Generator(device='cuda').manual_seed(seed)
+
+    #  get prompt embeds first
+    (
+    prompt_embeds,
+    negative_prompt_embeds,
+    pooled_prompt_embeds,
+    negative_pooled_prompt_embeds
+    ) = pipe.encode_prompt(prompt)
+    del pipe.text_encoder, pipe.text_encoder_2
+    flush()
     options = {
-        "prompt": [prompt],
-        "negative_prompt": [negative_prompt],
-        "negative_prompt_2": [neg_prompt_2],
+     #   "prompt": [prompt],
+      #  "negative_prompt": [negative_prompt],
+     #   "negative_prompt_2": [neg_prompt_2],
+		"prompt_embeds": = [prompt_embeds],
+		"negative_prompt_embeds": = [negative_prompt_embeds],
+		"pooled_prompt_embeds": = [pooled_prompt_embeds],
+		"negative_pooled_prompt_embeds": = [negative_pooled_prompt_embeds],
         "width": width,
         "height": height,
         "guidance_scale": guidance_scale,
@@ -266,10 +292,24 @@ def generate_90(
 ):
     seed = random.randint(0, MAX_SEED)
     generator = torch.Generator(device='cuda').manual_seed(seed)
+
+    #  get prompt embeds first
+    (
+    prompt_embeds,
+    negative_prompt_embeds,
+    pooled_prompt_embeds,
+    negative_pooled_prompt_embeds
+    ) = pipe.encode_prompt(prompt)
+    del pipe.text_encoder, pipe.text_encoder_2
+    flush()
     options = {
-        "prompt": [prompt],
-        "negative_prompt": [negative_prompt],
-        "negative_prompt_2": [neg_prompt_2],
+     #   "prompt": [prompt],
+      #  "negative_prompt": [negative_prompt],
+     #   "negative_prompt_2": [neg_prompt_2],
+		"prompt_embeds": = [prompt_embeds],
+		"negative_prompt_embeds": = [negative_prompt_embeds],
+		"pooled_prompt_embeds": = [pooled_prompt_embeds],
+		"negative_pooled_prompt_embeds": = [negative_pooled_prompt_embeds],
         "width": width,
         "height": height,
         "guidance_scale": guidance_scale,