Update app.py

app.py

@@ -106,10 +106,9 @@ def apply_style(style_name: str, positive: str, negative: str = "") -> Tuple[str
         negative = ""
     return p.replace("{prompt}", positive), n + negative
     
-unetX = UNet2DConditionModel.from_pretrained('ford442/RealVisXL_V5.0_BF16', subfolder='unet', low_cpu_mem_usage=False, token=True) #.to(device).to(torch.bfloat16) #.to(device=device, dtype=torch.bfloat16)
 
 def load_and_prepare_model():
-    #vae = AutoencoderKL.from_pretrained("ford442/sdxl-vae-bf16", safety_checker=None)
+    unetX = UNet2DConditionModel.from_pretrained('ford442/RealVisXL_V5.0_BF16', subfolder='unet', low_cpu_mem_usage=False, token=True) #.to(device).to(torch.bfloat16) #.to(device=device, dtype=torch.bfloat16)
     vaeX = AutoencoderKL.from_pretrained("stabilityai/sdxl-vae", safety_checker=None, use_safetensors=False, low_cpu_mem_usage=False, torch_dtype=torch.float32, token=True) #.to(device).to(torch.bfloat16) #.to(device=device, dtype=torch.bfloat16)
     pipe = StableDiffusionXLPipeline.from_pretrained(
         'ford442/RealVisXL_V5.0_BF16',
@@ -122,7 +121,7 @@ def load_and_prepare_model():
         text_encoder=None,
         text_encoder_2=None,
         vae=None,
-        #unet=None,
+        unet=unetX,
     )
     
     '''
@@ -176,7 +175,7 @@ model = Phi3ForCausalLM.from_pretrained(checkpoint).to('cuda:0')
 #model = AutoModelForCausalLM.from_pretrained(checkpoint, device_map='cuda') #.to('cuda')
 
 ip_model = IPAdapterXL(pipe, local_folder, ip_ckpt, device)
-text_encoder=CLIPTextModel.from_pretrained('ford442/RealVisXL_V5.0_BF16', subfolder='text_encoder',token=True) #.to(device=device, dtype=torch.bfloat16)
+text_encoder_1=CLIPTextModel.from_pretrained('ford442/RealVisXL_V5.0_BF16', subfolder='text_encoder',token=True) #.to(device=device, dtype=torch.bfloat16)
 text_encoder_2=CLIPTextModelWithProjection.from_pretrained('ford442/RealVisXL_V5.0_BF16', subfolder='text_encoder_2',token=True) #.to(device=device, dtype=torch.bfloat16)
    
 MAX_SEED = np.iinfo(np.int32).max
@@ -405,6 +404,9 @@ def generate_30(
         print(caption)
         print(caption_2)
         print("-- generating further caption --")
+        global model5
+        global captioner2
+        global processor5
         del captioner2
         del model5
         del processor5
@@ -415,13 +417,17 @@ def generate_30(
         expanded = expand_prompt(caption_2)
         expanded_1 = expanded[0]
         expanded_2 = expanded[1]
+        global model
+        global txt_tokenizer
         del model
         del txt_tokenizer
         gc.collect()
         torch.cuda.clear_cache()
-        pipe.text_encoder=text_encoder.to(device=device, dtype=torch.bfloat16)
+        global text_encoder_1
+        global text_encoder_2
+        pipe.text_encoder=text_encoder_1.to(device=device, dtype=torch.bfloat16)
         pipe.text_encoder_2=text_encoder_2.to(device=device, dtype=torch.bfloat16)
-        pipe.unet=unetX.to(device=device, dtype=torch.bfloat16)
+        #pipe.unet=unetX.to(device=device, dtype=torch.bfloat16)
         
         print('-- generating image --')
         sd_image = ip_model.generate(
@@ -485,7 +491,7 @@ def generate_60(
     samples=1,
     progress=gr.Progress(track_tqdm=True)  # Add progress as a keyword argument
 ):
-    pipe.text_encoder=text_encoder
+    pipe.text_encoder=text_encoder_1
     pipe.text_encoder_2=text_encoder_2
     seed = random.randint(0, MAX_SEED)
     generator = torch.Generator(device='cuda').manual_seed(seed)
@@ -575,7 +581,7 @@ def generate_90(
     samples=1,
     progress=gr.Progress(track_tqdm=True)  # Add progress as a keyword argument
 ):
-    pipe.text_encoder=text_encoder
+    pipe.text_encoder=text_encoder_1
     pipe.text_encoder_2=text_encoder_2
     seed = random.randint(0, MAX_SEED)
     generator = torch.Generator(device='cuda').manual_seed(seed)