use 16bit only if cuda

app.py

@@ -94,7 +94,7 @@ def main(
     generator = torch.Generator().manual_seed(int(seed)) # use cpu as does not work on mps
 
     embeddings = base64_to_embedding(embeddings)
-    embeddings = torch.tensor(embeddings).to(device)
+    embeddings = torch.tensor(embeddings, dtype=torch_size).to(device)
 
     images_list = pipe(
         # inp.tile(n_samples, 1, 1, 1),
@@ -178,12 +178,12 @@ def on_embeddings_changed_update_plot(embeddings_b64):
             # height=300,
             width=embeddings.shape[0])
 
-
 device = torch.device("mps" if torch.backends.mps.is_available() else "cuda:0" if torch.cuda.is_available() else "cpu")
+torch_size = torch.float16 if device == ('cuda') else torch.float32
 pipe = StableDiffusionPipeline.from_pretrained(
     model_id,
     custom_pipeline="pipeline.py",
-    torch_dtype=torch.float16,
+    torch_dtype=torch_size,
     # , revision="fp16",
     requires_safety_checker = False, safety_checker=None,
     text_encoder = CLIPTextModel,