Update app.py

app.py

@@ -77,7 +77,7 @@ def load_model(base_model_path, lora_path):
         base_model_path,
         transformer=transformer,
         torch_dtype=torch.bfloat16
-    ).to("cuda")
+    )
     pipe.transformer.to(torch.bfloat16)
     gr.Info(str(f"Model loading: {int((80 / 100) * 100)}%"))
     gr.Info(str(f"Inject LoRA: {lora_path}"))
@@ -148,6 +148,7 @@ def predict(
         gray_image_pil = Image.fromarray(gray_image).convert('L')
     else:
         gray_image_pil = input_image["layers"][0]
+    pipe.to("cuda")
     result = pipe(
         prompt=prompt,
         control_image=input_image["background"].convert("RGB"),