device fix

image_generator.py

@@ -7,8 +7,6 @@ from transformers import CLIPTextModel, CLIPTokenizer, logging
 
 from utils import load_embedding_bin, set_timesteps, latents_to_pil
 from loss import blue_loss, cosine_loss
-from matplotlib import pyplot as plt
-from pathlib import Path
 
 torch.manual_seed(11)
 logging.set_verbosity_error()
@@ -43,10 +41,9 @@ vae = AutoencoderKL.from_pretrained(
 #
 # # 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").to(
-    torch_device
-)
-#
+text_encoder = CLIPTextModel.from_pretrained(
+    "openai/clip-vit-large-patch14").to(torch_device)
+
 # # The UNet model for generating the latents.
 unet = UNet2DConditionModel.from_pretrained(
     "CompVis/stable-diffusion-v1-4", subfolder="unet"
@@ -60,9 +57,6 @@ scheduler = LMSDiscreteScheduler(
     num_train_timesteps=1000,
 )
 
-# vae = vae
-# text_encoder = text_encoder.to(torch_device)
-unet = unet
 token_emb_layer = text_encoder.text_model.embeddings.token_embedding
 pos_emb_layer = text_encoder.text_model.embeddings.position_embedding
 position_ids = text_encoder.text_model.embeddings.position_ids[:, :77]
@@ -227,7 +221,7 @@ def generate_image_from_embeddings(
         )
 
         #### ADDITIONAL GUIDANCE ###
-        if i % 2 == 0:
+        if i % 5 == 0:
             # Requires grad on the latents
             latents = latents.detach().requires_grad_()
 

utils.py

@@ -1,8 +1,19 @@
+import os
 import torch
 from PIL import Image
 from diffusers import AutoencoderKL
 
-vae = AutoencoderKL.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="vae").to("mps:0")
+torch_device = (
+    "cuda"
+    if torch.cuda.is_available()
+    else "mps"
+    if torch.backends.mps.is_available()
+    else "cpu"
+)
+if "mps" == torch_device:
+    os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
+
+vae = AutoencoderKL.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="vae").to(torch_device)
 
 def pil_to_latent(input_im):
     # Single image -> single latent in a batch (so size 1, 4, 64, 64)