Update utils.py

utils.py

@@ -27,24 +27,22 @@ import gc
 gc.collect()
 torch.cuda.empty_cache()
 
-# Load the autoencoder model which will be used to decode the latents into image space.
-vae = AutoencoderKL.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="vae")
+from diffusers import StableDiffusionPipeline
 
-# Load the tokenizer and text encoder to tokenize and encode the text.
-tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14")
-text_encoder = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14")
+model_id = "segmind/tiny-sd"
 
-# The UNet model for generating the latents.
-unet = UNet2DConditionModel.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="unet")
+pipe = StableDiffusionPipeline.from_pretrained(model_id).to("cpu")
+text_encoder = pipe.text_encoder.to(torch_device)
+text_encoder.eval()
+unet = pipe.unet.to(torch_device)
+unet.eval()
+vae = pipe.vae.to(torch_device)
+vae.eval()
 
-# The noise scheduler
-scheduler = LMSDiscreteScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear",
-                                 num_train_timesteps=1000)
-
-# To the GPU we go!
-vae = vae.to(torch_device)
-text_encoder = text_encoder.to(torch_device)
-unet = unet.to(torch_device)
+tokenizer = CLIPTokenizer.from_pretrained('openai/clip-vit-large-patch14')
+scheduler = LMSDiscreteScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000)
+del pipe
+gc.collect()
 
 seed_values = [0, 0, 0, 0, 0]