mrleo1nid/ACE-Step
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

asign cuda device

Browse Source
This commit is contained in:
chuxij
2025-05-14 03:37:07 +00:00
parent 6003d6177d
commit 292ccef7bf
1 changed files with 4 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
acestep/gui.py
+4 -5
View File
@@ -9,11 +9,6 @@ Apache 2.0 License
import os import os
import click import click
from acestep.ui.components import create_main_demo_ui
from acestep.pipeline_ace_step import ACEStepPipeline
from acestep.data_sampler import DataSampler
@click.command() @click.command()
@click.option( @click.option(
"--checkpoint_path", "--checkpoint_path",
@@ -59,6 +54,10 @@ def main(checkpoint_path, server_name, port, device_id, share, bf16, torch_compi
os.environ["CUDA_VISIBLE_DEVICES"] = str(device_id) os.environ["CUDA_VISIBLE_DEVICES"] = str(device_id)
from acestep.ui.components import create_main_demo_ui
from acestep.pipeline_ace_step import ACEStepPipeline
from acestep.data_sampler import DataSampler
model_demo = ACEStepPipeline( model_demo = ACEStepPipeline(
checkpoint_dir=checkpoint_path, checkpoint_dir=checkpoint_path,
dtype="bfloat16" if bf16 else "float32", dtype="bfloat16" if bf16 else "float32",