Update app.py

app.py

@@ -23,12 +23,12 @@ class ModelWrapper:
         self.DTYPE = torch.float16
         self.device = 0
 
-        self.tokenizer_one = AutoTokenizer.from_pretrained(model_id, subfolder="tokenizer", revision=revision, use_fast=False)
-        self.tokenizer_two = AutoTokenizer.from_pretrained(model_id, subfolder="tokenizer", revision=revision, use_fast=False)
+        self.tokenizer_one = AutoTokenizer.from_pretrained(model_id, revision=revision, use_fast=False)
+        self.tokenizer_two = AutoTokenizer.from_pretrained(model_id, revision=revision, use_fast=False)
 
         self.text_encoder = SDXLTextEncoder(model_id, revision, accelerator, dtype=self.DTYPE)
 
-        self.vae = AutoencoderKL.from_pretrained(model_id, subfolder="vae").float().to(self.device)
+        self.vae = AutoencoderKL.from_pretrained(model_id).float().to(self.device)
         self.vae_dtype = torch.float32
 
         self.tiny_vae = AutoencoderTiny.from_pretrained("madebyollin/taesdxl", torch_dtype=self.DTYPE).to(self.device)
@@ -43,12 +43,12 @@ class ModelWrapper:
         self.vae_downsample_ratio = image_resolution // latent_resolution
         self.conditioning_timestep = conditioning_timestep
 
-        self.scheduler = DDIMScheduler.from_pretrained(model_id, subfolder="scheduler")
+        self.scheduler = DDIMScheduler.from_pretrained(model_id)
         self.alphas_cumprod = self.scheduler.alphas_cumprod.to(self.device)
         self.num_step = num_step
     
     def create_generator(self, model_id, checkpoint_path):
-        generator = UNet2DConditionModel.from_pretrained(model_id, subfolder="unet").to(self.DTYPE)
+        generator = UNet2DConditionModel.from_pretrained(model_id).to(self.DTYPE)
         state_dict = torch.load(checkpoint_path)
         generator.load_state_dict(state_dict, strict=True)
         generator.requires_grad_(False)
@@ -80,7 +80,6 @@ class ModelWrapper:
     @spaces.GPU()
     def sample(self, noise, unet_added_conditions, prompt_embed, fast_vae_decode):
         alphas_cumprod = self.scheduler.alphas_cumprod.to(self.device)
-        print(alphas_cumprod)
         if self.num_step == 1:
             all_timesteps = [self.conditioning_timestep]
             step_interval = 0 
@@ -173,8 +172,8 @@ class SDXLTextEncoder(torch.nn.Module):
     def __init__(self, model_id, revision, accelerator, dtype=torch.float16):
         super().__init__()
 
-        self.text_encoder_one = CLIPTextModel.from_pretrained(model_id, subfolder="text_encoder", revision=revision).to(0).to(dtype=dtype)
-        self.text_encoder_two = CLIPTextModelWithProjection.from_pretrained(model_id, subfolder="text_encoder_2", revision=revision).to(0).to(dtype=dtype)
+        self.text_encoder_one = CLIPTextModel.from_pretrained(model_id, revision=revision).to(0).to(dtype=dtype)
+        self.text_encoder_two = CLIPTextModelWithProjection.from_pretrained(model_id, revision=revision).to(0).to(dtype=dtype)
 
         self.accelerator = accelerator