Update app.py

app.py

@@ -5,7 +5,7 @@ import random
 
 import torch
 from diffusers import StableDiffusion3Pipeline, AutoencoderKL
-#from transformers import CLIPTextModelWithProjection, T5EncoderModel
+from transformers import CLIPTextModelWithProjection, T5EncoderModel
 from transformers import CLIPTokenizer, T5TokenizerFast
 
 import re
@@ -61,27 +61,35 @@ def upload_to_ftp(filename):
 pyx = cyper.inline(code, fast_indexing=True, directives=dict(boundscheck=False, wraparound=False, language_level=3))
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-#vaeX=AutoencoderKL.from_pretrained("ford442/sdxl-vae-bf16", use_safetensors=False, subfolder='sd3-vae',token=True)
+vaeX=AutoencoderKL.from_pretrained("ford442/stable-diffusion-3.5-large-fp32", use_safetensors=True, subfolder='vae',token=True)
 
 pipe = StableDiffusion3Pipeline.from_pretrained(
     #"stabilityai  #  stable-diffusion-3.5-large",
     "ford442/stable-diffusion-3.5-large-bf16",
 #    vae=AutoencoderKL.from_pretrained("ford442/stable-diffusion-3.5-large-fp32", use_safetensors=True, subfolder='vae',token=True),
      #scheduler = FlowMatchHeunDiscreteScheduler.from_pretrained('ford442/stable-diffusion-3.5-large-bf16', subfolder='scheduler',token=True),
+     text_encoder=None, #CLIPTextModelWithProjection.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder', token=True),
     # text_encoder=CLIPTextModelWithProjection.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder', token=True),
+    text_encoder_2=None, #CLIPTextModelWithProjection.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder_2',token=True),
   #  text_encoder_2=CLIPTextModelWithProjection.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder_2',token=True),
+    text_encoder_3=None, #T5EncoderModel.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder_3',token=True),
   #  text_encoder_3=T5EncoderModel.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder_3',token=True),
     #tokenizer=CLIPTokenizer.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", add_prefix_space=True, subfolder="tokenizer", token=True),
     #tokenizer_2=CLIPTokenizer.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", add_prefix_space=True, subfolder="tokenizer_2", token=True),
     tokenizer_3=T5TokenizerFast.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", add_prefix_space=False, use_fast=True, subfolder="tokenizer_3", token=True),
-    #vae=vaeX,
+    vae=None,
     #torch_dtype=torch.bfloat16,
     #use_safetensors=False,
 )
 #pipe.load_lora_weights("ford442/sdxl-vae-bf16", weight_name="LoRA/UltraReal.safetensors")
 pipe.to(device=device, dtype=torch.bfloat16)
 #pipe.to(device)
+pipe.vae=vaeX.to('cpu')
 
+text_encoder=CLIPTextModelWithProjection.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder', token=True).to(device=device, dtype=torch.bfloat16)
+text_encoder_2=CLIPTextModelWithProjection.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder_2',token=True).to(device=device, dtype=torch.bfloat16)
+text_encoder_3=T5EncoderModel.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder_3',token=True).to(device=device, dtype=torch.bfloat16)
+    
 upscaler_2 = UpscaleWithModel.from_pretrained("Kim2091/ClearRealityV1").to(torch.device('cpu'))
 
 MAX_SEED = np.iinfo(np.int32).max
@@ -100,6 +108,9 @@ def infer_30(
     num_inference_steps,
     progress=gr.Progress(track_tqdm=True),
 ):
+    pipe.text_encoder=text_encoder
+    pipe.text_encoder_2=text_encoder_2
+    pipe.text_encoder_3=text_encoder_3
     torch.set_float32_matmul_precision("highest")
     seed = random.randint(0, MAX_SEED)
     generator = torch.Generator(device='cuda').manual_seed(seed)
@@ -147,6 +158,9 @@ def infer_60(
     num_inference_steps,
     progress=gr.Progress(track_tqdm=True),
 ):
+    pipe.text_encoder=text_encoder
+    pipe.text_encoder_2=text_encoder_2
+    pipe.text_encoder_3=text_encoder_3
     torch.set_float32_matmul_precision("highest")
     seed = random.randint(0, MAX_SEED)
     generator = torch.Generator(device='cuda').manual_seed(seed)
@@ -193,6 +207,9 @@ def infer_90(
     num_inference_steps,
     progress=gr.Progress(track_tqdm=True),
 ):
+    pipe.text_encoder=text_encoder
+    pipe.text_encoder_2=text_encoder_2
+    pipe.text_encoder_3=text_encoder_3
     torch.set_float32_matmul_precision("highest")
     seed = random.randint(0, MAX_SEED)
     generator = torch.Generator(device='cuda').manual_seed(seed)