Skip to content

Fix MPS device support for Apple Silicon Macs #26

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
btebedo wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Fix MPS device support for Apple Silicon Macs #26

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 8 additions & 2 deletions src/module/stable_diffusion.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,12 @@ def freeze(model):
class StableDiffusion(Base):
def setup(self):
hf_key = self.config.get("hf_key", None)
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
if torch.cuda.is_available():
self.device = torch.device("cuda")
elif torch.backends.mps.is_available():
self.device = torch.device("mps")
else:
self.device = torch.device("cpu")
fp16 = self.config.get("fp16", True)
self.dtype = torch.bfloat16 if fp16 else torch.float32
vae_padding = self.config.get("vae_padding", "zeros")
Expand Down Expand Up @@ -62,7 +67,7 @@ def setup(self):
# 3. Text encoder (CLIP)
text_encoder_config = CLIPTextConfig.from_pretrained(model_key, subfolder="text_encoder", local_files_only=local_files_only)
self.text_encoder = CLIPTextModel(text_encoder_config)
self.text_encoder.to(self.device, dtype=self.dtype).eval()
self.text_encoder = self.text_encoder.to(self.device, dtype=self.dtype).eval()
# 4. Tokenizer (CLIP tokenizer, this one has vocab so from_pretrained is needed)
self.tokenizer = CLIPTokenizer.from_pretrained(model_key, subfolder="tokenizer", local_files_only=local_files_only)
# 5. Scheduler
Expand All @@ -74,6 +79,7 @@ def setup(self):

def encode_text(self, prompt, padding_mode="do_not_pad"):
# prompt: [str]
self.text_encoder = self.text_encoder.to(self.device) # force materialization for MPS
inputs = self.tokenizer(
prompt,
padding=padding_mode,
Expand Down