Update app.py

app.py

@@ -4,8 +4,8 @@ import numpy as np
 #import tensorrt as trt
 import random
 import torch
-from diffusers import StableDiffusion3Pipeline, AutoencoderKL 
-from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
+from diffusers import StableDiffusion3Pipeline, AutoencoderKL, SD3Transformer2DModel
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer, CLIPTextModelWithProjection, T5EncoderModel
 #from threading import Thread
 #from transformers import pipeline
 from transformers import T5Tokenizer, T5ForConditionalGeneration
@@ -63,7 +63,13 @@ checkpoint = "microsoft/Phi-3.5-mini-instruct"
 #vaeXL = AutoencoderKL.from_pretrained("stabilityai/sdxl-vae", safety_checker=None, use_safetensors=False) #, device_map='cpu') #.to(torch.bfloat16) #.to(device=device, dtype=torch.bfloat16)
 
 pipe = StableDiffusion3Pipeline.from_pretrained(
-    "ford442/stable-diffusion-3.5-large-bf16",
+    "stabilityai/stable-diffusion-3.5-large",
+    #"ford442/stable-diffusion-3.5-large-bf16",
+    vae=AutoencoderKL.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='vae',token=True)
+    transformer=SD3Transformer2DModel.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='transformer',token=True)
+    text_encoder=CLIPTextModelWithProjection.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder',token=True),
+    text_encoder_2=CLIPTextModelWithProjection.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder_2',token=True),
+    text_encoder_3=T5EncoderModel.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder_3',token=True),
     token=True,
     use_safetensors=False,
 )