Update app.py

app.py

@@ -20,7 +20,6 @@ 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
@@ -86,18 +85,22 @@ 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")
+    
     pipe = StableDiffusionXLPipeline.from_pretrained(
         'ford442/RealVisXL_V5.0_BF16',
         #torch_dtype=torch.bfloat16,
         token=True,
-       # low_cpu_mem_usage = False,
         add_watermarker=False,
+        text_encoder=None,
+        text_encoder_2=None,
+        vae=None,
     )
+    #pipe.vae = vaeXL #.to(torch.bfloat16)
     #pipe.scheduler = sched
     #pipe.vae.do_resize=False
     #pipe.vae.vae_scale_factor=8
@@ -110,8 +113,9 @@ 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.vae = vaeXL #.to(torch.bfloat16)
+    pipe.unet.set_attn_processor(AttnProcessor2_0())
+    
+    pipe.vae = vaeXL.to('cpu') #.to(torch.bfloat16)
     return pipe
     
 pipe = load_and_prepare_model()
@@ -164,11 +168,6 @@ 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,
@@ -184,24 +183,12 @@ 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,negative_prompt)
-    del pipe.text_encoder, pipe.text_encoder_2
-    flush()
+    pipe.text_encoder=CLIPTextModel.from_config(pipe.config)
+    pipe.text_encoder_2=CLIPTextModelWithProjection.from_config(pipe.config)
     options = {
-     #   "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,
+        "prompt": [prompt],
+        "negative_prompt": [negative_prompt],
+        "negative_prompt_2": [neg_prompt_2],
         "width": width,
         "height": height,
         "guidance_scale": guidance_scale,
@@ -238,24 +225,12 @@ 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()
+    pipe.text_encoder=CLIPTextModel.from_config(pipe.config)
+    pipe.text_encoder_2=CLIPTextModelWithProjection.from_config(pipe.config)
     options = {
-     #   "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,
+        "prompt": [prompt],
+        "negative_prompt": [negative_prompt],
+        "negative_prompt_2": [neg_prompt_2],
         "width": width,
         "height": height,
         "guidance_scale": guidance_scale,
@@ -292,24 +267,12 @@ 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()
+    pipe.text_encoder=CLIPTextModel.from_config(pipe.config)
+    pipe.text_encoder_2=CLIPTextModelWithProjection.from_config(pipe.config)
     options = {
-     #   "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,
+        "prompt": [prompt],
+        "negative_prompt": [negative_prompt],
+        "negative_prompt_2": [neg_prompt_2],
         "width": width,
         "height": height,
         "guidance_scale": guidance_scale,