Update hy3dgen/text2image.py

hy3dgen/text2image.py

@@ -14,12 +14,22 @@
 
 import os
 import random
-
 import numpy as np
 import torch
 from diffusers import AutoPipelineForText2Image
 
 
+# -------------------- Device Auto-Selection --------------------
+def get_auto_device():
+    if torch.cuda.is_available():
+        return "cuda"
+    elif torch.backends.mps.is_available():  # macOS GPU
+        return "mps"
+    else:
+        return "cpu"
+
+
+# -------------------- Seed Helper --------------------
 def seed_everything(seed):
     random.seed(seed)
     np.random.seed(seed)
@@ -31,43 +41,70 @@ class HunyuanDiTPipeline:
     def __init__(
         self,
         model_path="Tencent-Hunyuan/HunyuanDiT-v1.1-Diffusers-Distilled",
-        device='cuda'
+        device=None
     ):
+        # Device auto-detect
+        if device is None:
+            device = get_auto_device()
+
         self.device = device
+
+        # float16 only works on CUDA / sometimes MPS, but not CPU
+        dtype = torch.float16 if device in ["cuda", "mps"] else torch.float32
+
         self.pipe = AutoPipelineForText2Image.from_pretrained(
             model_path,
-            torch_dtype=torch.float16,
+            torch_dtype=dtype,
             enable_pag=True,
             pag_applied_layers=["blocks.(16|17|18|19)"]
         ).to(device)
+
         self.pos_txt = ",白色背景,3D风格,最佳质量"
-        self.neg_txt = "文本,特写,裁剪,出框,最差质量,低质量,JPEG伪影,PGLY,重复,病态," \
-                       "残缺,多余的手指,变异的手,画得不好的手,画得不好的脸,变异,畸形,模糊,脱水,糟糕的解剖学," \
-                       "糟糕的比例,多余的肢体,克隆的脸,毁容,恶心的比例,畸形的肢体,缺失的手臂,缺失的腿," \
-                       "额外的手臂,额外的腿,融合的手指,手指太多,长脖子"
+        self.neg_txt = (
+            "文本,特写,裁剪,出框,最差质量,低质量,JPEG伪影,PGLY,重复,病态,"
+            "残缺,多余的手指,变异的手,画得不好的手,画得不好的脸,变异,畸形,模糊,脱水,糟糕的解剖学,"
+            "糟糕的比例,多余的肢体,克隆的脸,毁容,恶心的比例,畸形的肢体,缺失的手臂,缺失的腿,"
+            "额外的手臂,额外的腿,融合的手指,手指太多,长脖子"
+        )
 
+    # -------------------- Compile (optional) --------------------
     def compile(self):
-        # accelarate hunyuan-dit transformer,first inference will cost long time
-        torch.set_float32_matmul_precision('high')
-        self.pipe.transformer = torch.compile(self.pipe.transformer, fullgraph=True)
-        # self.pipe.vae.decode = torch.compile(self.pipe.vae.decode, fullgraph=True)
-        generator = torch.Generator(device=self.pipe.device)  # infer once for hot-start
-        out_img = self.pipe(
-            prompt='美少女战士',
-            negative_prompt='模糊',
+        # accelerate transformer — works only on CUDA; skip otherwise
+        if self.device == "cuda":
+            torch.set_float32_matmul_precision("high")
+            self.pipe.transformer = torch.compile(self.pipe.transformer, fullgraph=True)
+
+        # Safe generator creation (mps can't use device=...)
+        try:
+            generator = torch.Generator(device=self.pipe.device)
+        except:
+            generator = torch.Generator()
+
+        # warmup inference
+        _ = self.pipe(
+            prompt="美少女战士",
+            negative_prompt="模糊",
             num_inference_steps=25,
             pag_scale=1.3,
             width=1024,
             height=1024,
-            generator=generator,
+            generator=generator.manual_seed(42),
             return_dict=False
         )[0][0]
 
+    # -------------------- Generate Image --------------------
     @torch.no_grad()
     def __call__(self, prompt, seed=0):
         seed_everything(seed)
-        generator = torch.Generator(device=self.pipe.device)
+
+        # Generator fix (no device for mps)
+        try:
+            generator = torch.Generator(device=self.pipe.device)
+        except:
+            generator = torch.Generator()
+
         generator = generator.manual_seed(int(seed))
+
         out_img = self.pipe(
             prompt=prompt[:60] + self.pos_txt,
             negative_prompt=self.neg_txt,
@@ -78,4 +115,5 @@ class HunyuanDiTPipeline:
             generator=generator,
             return_dict=False
         )[0][0]
+
         return out_img