Update app.py

app.py

@@ -112,15 +112,16 @@ 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('cpu') #.to(torch.bfloat16)
+    pipe.vae = vaeXL.to(device) #.to('cpu') #.to(torch.bfloat16)
+    
     pipe.unet.set_attn_processor(AttnProcessor2_0())
     pipe.vae.set_default_attn_processor()
     return pipe
     
 pipe = load_and_prepare_model()
 
-text_encoder=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)
+text_encoder=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
 
@@ -185,8 +186,8 @@ def generate_30(
 ):
     seed = random.randint(0, MAX_SEED)
     generator = torch.Generator(device='cuda').manual_seed(seed)
-    pipe.text_encoder=text_encoder
-    pipe.text_encoder_2=text_encoder_2
+    pipe.text_encoder=text_encoder.to(device=device, dtype=torch.bfloat16)
+    pipe.text_encoder_2=text_encoder_2.to(device=device, dtype=torch.bfloat16)
     options = {
         "prompt": [prompt],
         "negative_prompt": [negative_prompt],
@@ -227,8 +228,8 @@ def generate_60(
 ):
     seed = random.randint(0, MAX_SEED)
     generator = torch.Generator(device='cuda').manual_seed(seed)
-    pipe.text_encoder=text_encoder
-    pipe.text_encoder_2=text_encoder_2
+    pipe.text_encoder=text_encoder.to(device=device, dtype=torch.bfloat16)
+    pipe.text_encoder_2=text_encoder_2.to(device=device, dtype=torch.bfloat16)
     options = {
         "prompt": [prompt],
         "negative_prompt": [negative_prompt],
@@ -269,8 +270,8 @@ def generate_90(
 ):
     seed = random.randint(0, MAX_SEED)
     generator = torch.Generator(device='cuda').manual_seed(seed)
-    pipe.text_encoder=text_encoder
-    pipe.text_encoder_2=text_encoder_2
+    pipe.text_encoder=text_encoder.to(device=device, dtype=torch.bfloat16)
+    pipe.text_encoder_2=text_encoder_2.to(device=device, dtype=torch.bfloat16)
     options = {
         "prompt": [prompt],
         "negative_prompt": [negative_prompt],