Update app.py

app.py

@@ -1,45 +1,108 @@
+# app.py
 import argparse
-from ui.components import create_main_demo_ui
+import streamlit as st
+import os
 from pipeline_ace_step import ACEStepPipeline
 from data_sampler import DataSampler
-import os
-
-
-parser = argparse.ArgumentParser()
-parser.add_argument("--checkpoint_path", type=str, default=None)
-parser.add_argument("--server_name", type=str, default="0.0.0.0")
-parser.add_argument("--port", type=int, default=7860)
-parser.add_argument("--device_id", type=int, default=0)
-parser.add_argument("--share", action='store_true', default=False)
-parser.add_argument("--bf16", action='store_true', default=True)
-parser.add_argument("--torch_compile", type=bool, default=False)
-
-args = parser.parse_args()
-os.environ["CUDA_VISIBLE_DEVICES"] = str(args.device_id)
-
 
-persistent_storage_path = "/data"
+# 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'
+    }
 
-def main(args):
-
+@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",
+        checkpoint_dir=args['checkpoint_path'],
+        dtype="bfloat16" if args['bf16'] else "float32",
         persistent_storage_path=persistent_storage_path,
-        torch_compile=args.torch_compile
+        torch_compile=args['torch_compile']
     )
     data_sampler = DataSampler()
+    return model_demo, data_sampler
 
-    demo = create_main_demo_ui(
-        text2music_process_func=model_demo.__call__,
-        sample_data_func=data_sampler.sample,
-        load_data_func=data_sampler.load_json,
-    )
-    demo.queue(default_concurrency_limit=8).launch(
+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(args)
+    main()