force cuda set

app.py

@@ -122,6 +122,11 @@ class ImageGenerator:
         self.ae = self.ae.to(device=self.device, dtype=torch.float32)
         self.dit = self.dit.to(device=self.device, dtype=dtype)
         self.llm_encoder = self.llm_encoder.to(device=self.device, dtype=dtype)
+    
+    def to_cuda(self):
+        self.ae.to(device='cuda', dtype=torch.float32)
+        self.dit.to(device='cuda', dtype=torch.bfloat16)
+        self.llm_encoder.to(device='cuda', dtype=torch.bfloat16)
 
     def prepare(self, prompt, img, ref_image, ref_image_raw):
         bs, _, h, w = img.shape
@@ -377,10 +382,32 @@ class ImageGenerator:
         return images_list
 
 
+# 模型仓库ID（如："bert-base-uncased"）
+model_repo = "stepfun-ai/Step1X-Edit"
+# 本地保存路径
+model_path = "./model_weights"
+os.makedirs(model_path, exist_ok=True)
+
+
+# 下载模型（包括所有文件）
+snapshot_download(
+    repo_id=model_repo,
+    local_dir=model_path,
+    local_dir_use_symlinks=False  # 避免使用符号链接
+)
+
+
+image_edit = ImageGenerator(
+    ae_path=os.path.join(model_path, 'vae.safetensors'),
+    dit_path=os.path.join(model_path, "step1x-edit-i1258.safetensors"),
+    qwen2vl_model_path='Qwen/Qwen2.5-VL-7B-Instruct',
+    max_length=640,
+)
+
 
 
 @spaces.GPU(duration=240)
-def inference(prompt, ref_images, seed, size_level, infer_func=None):
+def inference(prompt, ref_images, seed, size_level):
     start_time = time.time()
 
     if seed == -1:
@@ -389,7 +416,11 @@ def inference(prompt, ref_images, seed, size_level, infer_func=None):
     else:
         random_seed = seed
 
-    image = infer_func(
+    image_edit.to_cuda()
+
+    inference_func = image_edit.generate_image
+    
+    image = inference_func(
         prompt,
         negative_prompt="",
         ref_images=ref_images.convert('RGB'),
@@ -404,32 +435,6 @@ def inference(prompt, ref_images, seed, size_level, infer_func=None):
     print(f"Time taken: {time.time() - start_time:.2f} seconds")
     return image, random_seed
 
-
-# 模型仓库ID（如："bert-base-uncased"）
-model_repo = "stepfun-ai/Step1X-Edit"
-# 本地保存路径
-model_path = "./model_weights"
-os.makedirs(model_path, exist_ok=True)
-
-
-# 下载模型（包括所有文件）
-snapshot_download(
-    repo_id=model_repo,
-    local_dir=model_path,
-    local_dir_use_symlinks=False  # 避免使用符号链接
-)
-
-
-image_edit = ImageGenerator(
-    ae_path=os.path.join(model_path, 'vae.safetensors'),
-    dit_path=os.path.join(model_path, "step1x-edit-i1258.safetensors"),
-    qwen2vl_model_path='Qwen/Qwen2.5-VL-7B-Instruct',
-    max_length=640,
-)
-
-inference_func = image_edit.generate_image
-
-    # inference_func = prepare_infer_func()
 with gr.Blocks() as demo:
     gr.Markdown(
         """