app.py

# app.py
import argparse
import streamlit as st
import os
from pipeline_ace_step import ACEStepPipeline
from data_sampler import DataSampler

# Streamlit 설정
st.set_page_config(
    page_title="ACE Step Music Generator",
    page_icon="🎵",
    layout="wide"
)

def get_args():
    """환경변수 또는 기본값으로 설정"""
    return {
        'checkpoint_path': os.environ.get('CHECKPOINT_PATH'),
        'device_id': int(os.environ.get('DEVICE_ID', '0')),
        'bf16': os.environ.get('BF16', 'True').lower() == 'true',
        'torch_compile': os.environ.get('TORCH_COMPILE', 'False').lower() == 'true'
    }

@st.cache_resource
def load_model(args):
    """모델 로딩 (캐시됨)"""
    os.environ["CUDA_VISIBLE_DEVICES"] = str(args['device_id'])
    persistent_storage_path = "/data"
    
    model_demo = ACEStepPipeline(
        checkpoint_dir=args['checkpoint_path'],
        dtype="bfloat16" if args['bf16'] else "float32",
        persistent_storage_path=persistent_storage_path,
        torch_compile=args['torch_compile']
    )
    data_sampler = DataSampler()
    return model_demo, data_sampler

def main():
    st.title("🎵 ACE Step Music Generator")
    
    args = get_args()
    
    try:
        model_demo, data_sampler = load_model(args)
        
        # UI 구성
        col1, col2 = st.columns([2, 1])
        
        with col1:
            st.header("Generate Music")
            
            # 텍스트 입력
            prompt = st.text_area(
                "Enter your music description:",
                placeholder="Enter a description of the music you want to generate...",
                height=100
            )
            
            # 생성 버튼
            if st.button("Generate Music", type="primary"):
                if prompt:
                    with st.spinner("Generating music..."):
                        try:
                            result = model_demo(prompt)
                            st.success("Music generated successfully!")
                            
                            # 결과 표시 (result 형태에 따라 조정 필요)
                            if hasattr(result, 'audio'):
                                st.audio(result.audio)
                            else:
                                st.write(result)
                                
                        except Exception as e:
                            st.error(f"Error generating music: {str(e)}")
                else:
                    st.warning("Please enter a description first.")
        
        with col2:
            st.header("Sample Data")
            
            if st.button("Load Sample"):
                try:
                    sample_data = data_sampler.sample()
                    st.json(sample_data)
                except Exception as e:
                    st.error(f"Error loading sample: {str(e)}")
            
            # 파일 업로드
            uploaded_file = st.file_uploader(
                "Upload JSON data",
                type=['json']
            )
            
            if uploaded_file:
                try:
                    data = data_sampler.load_json(uploaded_file)
                    st.json(data)
                except Exception as e:
                    st.error(f"Error loading file: {str(e)}")
                    
    except Exception as e:
        st.error(f"Error loading model: {str(e)}")
        import traceback
        st.code(traceback.format_exc())

if __name__ == "__main__":
    main()