try: import spaces USING_SPACES = True except ImportError: USING_SPACES = False import zero import gradio as gr import sys import threading import queue from io import TextIOBase import datetime import subprocess import os from inference import postprocess_inst_names from inference import inference_patch from convert import abc2xml, xml2, pdf2img def gpu_decorator(func): if USING_SPACES: return spaces.GPU(func) else: return func # 读取 prompt 组合 with open('prompts.txt', 'r') as f: prompts = f.readlines() valid_combinations = set() for prompt in prompts: prompt = prompt.strip() parts = prompt.split('_') valid_combinations.add((parts[0], parts[1], parts[2])) # 准备下拉框选项 periods = sorted({p for p, _, _ in valid_combinations}) composers = sorted({c for _, c, _ in valid_combinations}) instruments = sorted({i for _, _, i in valid_combinations}) # 动态更新作曲家、乐器下拉选项 def update_components(period, composer): if not period: return [ gr.update(choices=[], value=None, interactive=False), gr.update(choices=[], value=None, interactive=False) ] valid_composers = sorted({c for p, c, _ in valid_combinations if p == period}) valid_instruments = sorted({i for p, c, i in valid_combinations if p == period and c == composer}) if composer else [] return [ gr.update( choices=valid_composers, value=composer if composer in valid_composers else None, interactive=True ), gr.update( choices=valid_instruments, value=None, interactive=bool(valid_instruments) ) ] # 自定义实时流，用于把模型推理过程输出到前端 class RealtimeStream(TextIOBase): def __init__(self, queue): self.queue = queue def write(self, text): self.queue.put(text) return len(text) def convert_files(abc_content, period, composer, instrumentation): if not all([period, composer, instrumentation]): raise gr.Error("Please complete a valid generation first before saving") timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") prompt_str = f"{period}_{composer}_{instrumentation}" filename_base = f"{timestamp}_{prompt_str}" abc_filename = f"{filename_base}.abc" with open(abc_filename, "w", encoding="utf-8") as f: f.write(abc_content) # instrumentation replacement postprocessed_inst_abc = postprocess_inst_names(abc_content) filename_base_postinst = f"{filename_base}_postinst" with open(filename_base_postinst + ".abc", "w", encoding="utf-8") as f: f.write(postprocessed_inst_abc) # 转换文件 file_paths = {'abc': abc_filename} try: # abc2xml abc2xml(filename_base) abc2xml(filename_base_postinst) # xml2pdf xml2(filename_base, 'pdf') # xml2mid xml2(filename_base, 'mid') xml2(filename_base_postinst, 'mid') # xml2wav xml2(filename_base, 'wav') xml2(filename_base_postinst, 'wav') # 将PDF转为图片 images = pdf2img(filename_base) for i, image in enumerate(images): image.save(f"{filename_base}_page_{i+1}.png", "PNG") file_paths.update({ 'xml': f"{filename_base_postinst}.xml", 'pdf': f"{filename_base}.pdf", 'mid': f"{filename_base_postinst}.mid", 'wav': f"{filename_base_postinst}.wav", 'pages': len(images), 'current_page': 0, 'base': filename_base }) except Exception as e: raise gr.Error(f"文件处理失败: {str(e)}") return file_paths # 翻页控制函数 def update_page(direction, data): """ data 里面包含了 'pages','current_page','base' 三个关键信息 """ if not data: return None, gr.update(interactive=False), gr.update(interactive=False), data if direction == "prev" and data['current_page'] > 0: data['current_page'] -= 1 elif direction == "next" and data['current_page'] < data['pages'] - 1: data['current_page'] += 1 current_page_index = data['current_page'] # 更新图片路径 new_image = f"{data['base']}_page_{current_page_index+1}.png" # 当 current_page==0 时，prev_btn 不可用；当 current_page==pages-1 时，next_btn 不可用 prev_btn_state = gr.update(interactive=(current_page_index > 0)) next_btn_state = gr.update(interactive=(current_page_index < data['pages'] - 1)) return new_image, prev_btn_state, next_btn_state, data @gpu_decorator def generate_music(period, composer, instrumentation): """ 需要保证每次 yield 的返回值数量一致。 我们这里准备返回 5 个值，对应: 1) process_output (中间推理信息) 2) final_output (最终 ABC) 3) pdf_image (PDF 第一页对应的 png 路径) 4) audio_player (WAV 路径) 5) pdf_state (翻页用的 state) """ if (period, composer, instrumentation) not in valid_combinations: # 如果组合非法，直接抛出错误 raise gr.Error("Invalid prompt combination! Please re-select from the period options") # # Ensure model weights were downloaded successfully # if not os.path.exists(model_weights_path): # raise gr.Error(f"Model weights not available at {model_weights_path}") output_queue = queue.Queue() original_stdout = sys.stdout sys.stdout = RealtimeStream(output_queue) result_container = [] def run_inference(): try: # 使用下载的模型权重路径进行推理 result = inference_patch(period, composer, instrumentation) result_container.append(result) finally: sys.stdout = original_stdout thread = threading.Thread(target=run_inference) thread.start() process_output = "" final_output_abc = "" pdf_image = None audio_file = None pdf_state = None # 先持续读中间输出 while thread.is_alive(): try: text = output_queue.get(timeout=0.1) process_output += text # 暂时没有最终 ABC，还没有转文件 yield process_output, final_output_abc, pdf_image, audio_file, pdf_state except queue.Empty: continue # 线程结束后，把剩余的队列都拿出来 while not output_queue.empty(): text = output_queue.get() process_output += text # 最终推理结果 final_result = result_container[0] if result_container else "" # 显示转换文件的提示 final_output_abc = "Converting files..." yield process_output, final_output_abc, pdf_image, audio_file, pdf_state # 做文件转换 try: file_paths = convert_files(final_result, period, composer, instrumentation) final_output_abc = final_result # 拿到第一张图片和 wav 文件 if file_paths['pages'] > 0: pdf_image = f"{file_paths['base']}_page_1.png" audio_file = file_paths['wav'] pdf_state = file_paths # 直接把转换后的信息字典拿来存到 state except Exception as e: # 如果失败了，把错误信息返回到输出框 yield process_output, f"Error converting files: {str(e)}", None, None, None return # 最后一次 yield，带上所有信息 yield process_output, final_output_abc, pdf_image, audio_file, pdf_state def get_file(file_type, period, composer, instrumentation): """ 返回本地的指定类型文件，用于 Gradio 下载 """ # 这里其实需要你根据先前保存下来的具体文件路径来返回，演示时可以简化 # 如果是按 timestamp 去匹配，可以把转换的文件都存在某个目录下再拿最新的 # 这里仅做示例: possible_files = [f for f in os.listdir('.') if f.endswith(f'.{file_type}')] if not possible_files: return None # 简单返回最新的 possible_files.sort(key=os.path.getmtime) return possible_files[-1] css = """ /* 紧凑按钮样式 */ button[size="sm"] { padding: 4px 8px !important; margin: 2px !important; min-width: 60px; } /* PDF预览区 */ #pdf-preview { border-radius: 8px; /* 圆角 */ box-shadow: 0 2px 8px rgba(0,0,0,0.1); /* 阴影 */ } .page-btn { padding: 12px !important; /* 增大点击区域 */ margin: auto !important; /* 垂直居中 */ } /* 按钮悬停效果 */ .page-btn:hover { background: #f0f0f0 !important; transform: scale(1.05); } /* 布局调整 */ .gr-row { gap: 10px !important; /* 元素间距 */ } /* 音频播放器 */ .audio-panel { margin-top: 15px !important; max-width: 400px; } #audio-preview audio { height: 200px !important; } /* 保存功能区 */ .save-as-row { margin-top: 15px; padding: 10px; border-top: 1px solid #eee; } .save-as-label { font-weight: bold; margin-right: 10px; align-self: center; } .save-buttons { gap: 5px; /* 按钮间距 */ } """ with gr.Blocks(css=css) as demo: gr.Markdown("## NotaGen") # 用于保存 PDF 页数、当前页等信息 pdf_state = gr.State() with gr.Column(): with gr.Row(): # 左侧栏 with gr.Column(): with gr.Row(): period_dd = gr.Dropdown( choices=periods, value=None, label="Period", interactive=True ) composer_dd = gr.Dropdown( choices=[], value=None, label="Composer", interactive=False ) instrument_dd = gr.Dropdown( choices=[], value=None, label="Instrumentation", interactive=False ) generate_btn = gr.Button("Generate!", variant="primary") process_output = gr.Textbox( label="Generation process", interactive=False, lines=2, max_lines=2, placeholder="Generation progress will be shown here..." ) final_output = gr.Textbox( label="Post-processed ABC notation scores", interactive=True, lines=8, max_lines=8, placeholder="Post-processed ABC scores will be shown here..." ) # 音频播放 audio_player = gr.Audio( label="Audio Preview", format="wav", interactive=False, # container=False, # elem_id="audio-preview" ) # 右侧栏 with gr.Column(): # 图片容器 pdf_image = gr.Image( label="Sheet Music Preview", show_label=False, height=650, type="filepath", elem_id="pdf-preview", interactive=False, show_download_button=False ) # 翻页按钮 with gr.Row(): prev_btn = gr.Button( "⬅️ Last Page", variant="secondary", size="sm", elem_classes="page-btn" ) next_btn = gr.Button( "Next Page ➡️", variant="secondary", size="sm", elem_classes="page-btn" ) # 按钮组 with gr.Row(): gr.Markdown("**Save As: (Scroll down to get the link)**") save_abc = gr.Button("🅰️ ABC", variant="secondary", size="sm") save_xml = gr.Button("🎼 XML", variant="secondary", size="sm") save_pdf = gr.Button("📑 PDF", variant="secondary", size="sm") save_mid = gr.Button("🎹 MIDI", variant="secondary", size="sm") save_wav = gr.Button("🎧 WAV", variant="secondary", size="sm") # save_status = gr.Textbox( # label="Save Status", # interactive=False, # visible=True, # max_lines=1 # ) # 下拉框联动 period_dd.change( update_components, inputs=[period_dd, composer_dd], outputs=[composer_dd, instrument_dd] ) composer_dd.change( update_components, inputs=[period_dd, composer_dd], outputs=[composer_dd, instrument_dd] ) # 点击生成按钮，注意 outputs 要和 generate_music 里每次 yield 保持一致 generate_btn.click( generate_music, inputs=[period_dd, composer_dd, instrument_dd], outputs=[process_output, final_output, pdf_image, audio_player, pdf_state] ) # 翻页 prev_signal = gr.Textbox(value="prev", visible=False) next_signal = gr.Textbox(value="next", visible=False) prev_btn.click( update_page, inputs=[prev_signal, pdf_state], # ✅ 使用组件 outputs=[pdf_image, prev_btn, next_btn, pdf_state] ) next_btn.click( update_page, inputs=[next_signal, pdf_state], # ✅ 使用组件 outputs=[pdf_image, prev_btn, next_btn, pdf_state] ) # 文件保存按钮 save_abc.click( lambda state: state.get('abc') if state else None, inputs=[pdf_state], outputs=gr.File(label="abc", visible=True) ) save_xml.click( lambda state: state.get('xml') if state else None, inputs=[pdf_state], outputs=gr.File(label="xml", visible=True) ) save_pdf.click( lambda state: state.get('pdf') if state else None, inputs=[pdf_state], outputs=gr.File(label="pdf", visible=True) ) save_mid.click( lambda state: state.get('mid') if state else None, inputs=[pdf_state], outputs=gr.File(label="midi", visible=True) ) save_wav.click( lambda state: state.get('wav') if state else None, inputs=[pdf_state], outputs=gr.File(label="wav", visible=True) ) if __name__ == "__main__": # Configure GPU/CPU handling import torch # Function to initialize CUDA safely and verify it's working def is_cuda_working(): try: if torch.cuda.is_available(): # Test CUDA initialization with a small operation test_tensor = torch.tensor([1.0], device="cuda") _ = test_tensor * 2 return True return False except Exception as e: print(f"CUDA initialization test failed: {e}") return False # Check if running on Hugging Face Spaces if "SPACE_ID" in os.environ: cuda_working = is_cuda_working() if cuda_working: print("GPU is available and working. Using CUDA.") # You might want to set some environment variables or configurations here os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128" else: print("CUDA not working properly. Forcing CPU mode.") os.environ["CUDA_VISIBLE_DEVICES"] = "" torch.backends.cudnn.enabled = False # Launch with minimal parameters on Spaces demo.launch() else: # Running locally - use custom server settings and share print(f"Running locally with device: {'cuda' if torch.cuda.is_available() else 'cpu'}") demo.launch( server_name="0.0.0.0", server_port=7860, share=True # 确保外部访问 )