starred/ACE-Step
1
0
Fork 0
mirror of https://github.com/outbackdingo/ACE-Step.git synced 2026-03-20 22:46:08 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #167 from ace-step/set_cuda_device

Browse Source 
asign cuda device
This commit is contained in:
Gong Junmin
2025-05-14 11:39:52 +08:00
committed by GitHub
parent 9e4cc410d3 292ccef7bf
commit 39043aaa61
1 changed files with 4 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
acestep/gui.py
View File

@@ -9,11 +9,6 @@ Apache 2.0 License
import os
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.option(
"--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)
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(
checkpoint_dir=checkpoint_path,
dtype="bfloat16" if bf16 else "float32",