Update app.py

app.py

@@ -1,104 +1,131 @@
 import gradio as gr
-from transformers import pipeline
 import requests
-
-# 初始化模型（使用较简单的初始化方式避免内存问题）
-def load_model():
-    try:
-        # 尝试加载模型，如果内存不足则使用API
-        protein_pipeline = pipeline(
-            "text-generation", 
-            model="mzcwd/ProtTeX",
-            device=0 if torch.cuda.is_available() else -1
-        )
-        return protein_pipeline
-    except Exception as e:
-        print(f"模型加载失败: {e}")
-        return None
-
-model = load_model()
+import time
 
 def generate_protein(instruction, max_length=100):
     """生成蛋白质序列"""
     
-    # 如果本地模型加载失败，使用HuggingFace API
-    if model is None:
-        try:
-            API_URL = "https://api-inference.huggingface.co/models/mzcwd/ProtTeX"
-            headers = {"Authorization": f"Bearer {os.environ.get('ProtTex', '')}"}  # 需要替换成您的token
-            
-            response = requests.post(API_URL, json={
-                "inputs": instruction,
-                "parameters": {"max_length": max_length}
-            })
-            
-            if response.status_code == 200:
-                result = response.json()
-                return result[0]['generated_text']
-            else:
-                return f"API调用失败: {response.status_code}"
-                
-        except Exception as e:
-            return f"错误: {str(e)}"
+    # 显示加载状态
+    yield "正在生成蛋白质序列，请稍候..."
     
-    # 使用本地模型
     try:
-        result = model(
-            instruction, 
-            max_length=max_length,
-            num_return_sequences=1,
-            temperature=0.7
-        )
-        return result[0]['generated_text']
+        API_URL = "https://api-inference.huggingface.co/models/mzcwd/ProtTeX"
+        
+        payload = {
+            "inputs": instruction,
+            "parameters": {
+                "max_length": max_length,
+                "temperature": 0.8,
+                "do_sample": True,
+                "return_full_text": False
+            },
+            "options": {
+                "wait_for_model": True
+            }
+        }
+        
+        # 发送请求到 HuggingFace 推理 API
+        response = requests.post(API_URL, json=payload)
+        
+        if response.status_code == 200:
+            result = response.json()
+            if isinstance(result, list) and len(result) > 0:
+                generated_text = result[0].get('generated_text', '')
+                yield generated_text if generated_text else "模型已响应但未生成内容"
+            else:
+                yield f"响应格式异常: {result}"
+        elif response.status_code == 503:
+            # 模型正在加载，等待后重试
+            yield "模型正在启动中，请等待约30秒后重试..."
+        else:
+            error_msg = f"请求失败 (状态码: {response.status_code})"
+            if response.status_code == 401:
+                error_msg += "\n可能需要设置访问令牌"
+            yield error_msg
+            
     except Exception as e:
-        return f"生成失败: {str(e)}"
-
-# 示例指令
-examples = [
-    "Generate a protein sequence with alpha-helical structure",
-    "Create a beta-sheet rich protein for structural stability",
-    "Design a hydrophilic protein for aqueous environments",
-    "Generate a transmembrane protein sequence"
-]
+        yield f"发生错误: {str(e)}"
 
 # 创建界面
-with gr.Blocks(title="ProtTeX 蛋白质生成器") as demo:
-    gr.Markdown("# 🧬 ProtTeX 蛋白质生成器")
-    gr.Markdown("输入自然语言指令，生成蛋白质序列")
+with gr.Blocks(
+    title="ProtTeX 蛋白质生成器",
+    theme=gr.themes.Soft(),
+    css="""
+    .gradio-container {
+        max-width: 800px !important;
+    }
+    """
+) as demo:
+    
+    gr.Markdown("""
+    # 🧬 ProtTeX 蛋白质生成器
+    
+    使用自然语言指令生成蛋白质序列。输入您想要的蛋白质特性描述，AI将生成相应的蛋白质序列。
+    """)
     
     with gr.Row():
-        with gr.Column():
+        with gr.Column(scale=1):
             instruction = gr.Textbox(
-                label="输入指令",
-                placeholder="例如：Generate a protein sequence with alpha-helical structure",
-                lines=3
+                label="蛋白质生成指令",
+                placeholder="例如：Generate a protein with alpha-helical structure for membrane binding",
+                lines=3,
+                max_lines=5
             )
+            
             max_length = gr.Slider(
-                minimum=50, 
-                maximum=200, 
-                value=100, 
-                label="生成长度"
+                minimum=50,
+                maximum=300,
+                value=150,
+                step=10,
+                label="序列最大长度"
             )
-            generate_btn = gr.Button("生成蛋白质���列", variant="primary")
-        
-        with gr.Column():
+            
+            generate_btn = gr.Button(
+                "🚀 生成蛋白质序列",
+                variant="primary",
+                size="lg"
+            )
+            
+        with gr.Column(scale=1):
             output = gr.Textbox(
-                label="生成的序列",
-                lines=5,
+                label="生成的蛋白质序列",
+                lines=6,
+                show_copy_button=True,
                 interactive=False
             )
     
-    gr.Markdown("## 示例指令")
-    gr.Examples(
-        examples=examples,
-        inputs=instruction
+    # 示例部分
+    gr.Markdown("### 💡 示例指令")
+    
+    examples = gr.Examples(
+        examples=[
+            ["Generate a hydrophobic transmembrane protein sequence"],
+            ["Create a water-soluble protein with beta-sheet structure"],
+            ["Design a protein with enzymatic activity for hydrolysis"],
+            ["Generate a stable protein for high temperature environments"]
+        ],
+        inputs=[instruction],
+        outputs=[output],
+        fn=generate_protein,
+        cache_examples=False,
+        label="点击示例快速尝试"
     )
     
+    # 连接按钮事件
     generate_btn.click(
         fn=generate_protein,
         inputs=[instruction, max_length],
-        outputs=output
+        outputs=[output]
     )
+    
+    # 页脚信息
+    gr.Markdown("""
+    ---
+    *基于 [mzcwd/ProtTeX](https://huggingface.co/mzcwd/ProtTeX) 模型*
+    """)
 
 if __name__ == "__main__":
-    demo.launch(debug=True, share=True)
+    demo.launch(
+        server_name="0.0.0.0",
+        share=False
+    )