Update app.py

app.py

@@ -51,9 +51,9 @@ def load_dtr_from_huggingface():
         print(f"Error loading DTR dataset: {e}")
         return []
 
-def generate_natural_language_description(tags, api_key=None):
+def generate_natural_language_description(tags, api_key=None, base_url=None, model="gpt-4o"):
     """
-    使用 OpenAI GPT 生成自然语言描述。
+    使用 OpenAI GPT 或 DeepSeek API 生成自然语言描述。
     """
     if not api_key:
         api_key = os.getenv("OPENAI_API_KEY")
@@ -63,7 +63,7 @@ def generate_natural_language_description(tags, api_key=None):
     tag_descriptions = "\n".join([f"{key}: {value}" for key, value in tags.items() if value])
 
     try:
-        client = OpenAI(api_key=api_key)
+        client = OpenAI(api_key=api_key, base_url=base_url) if base_url else OpenAI(api_key=api_key)
 
         response = client.chat.completions.create(
             messages=[
@@ -80,7 +80,7 @@ def generate_natural_language_description(tags, api_key=None):
                     "content": f"Here are the tags and details:\n{tag_descriptions}\nPlease generate a vivid, imaginative scene description.",
                 },
             ],
-            model="gpt-4o",
+            model=model,
         )
         return response.choices[0].message.content.strip()
     except Exception as e:
@@ -88,7 +88,8 @@ def generate_natural_language_description(tags, api_key=None):
 
 def generate_prompt(
     action_file, style_file, artist_files, character_files, dtr_enabled, api_key, selected_categories,
-    expression_count, item_count, detail_count, scene_count, angle_count, quality_count, action_count, style_count
+    expression_count, item_count, detail_count, scene_count, angle_count, quality_count, action_count, style_count,
+    use_deepseek, deepseek_key
 ):
     """
     生成随机提示词和描述。
@@ -116,7 +117,10 @@ def generate_prompt(
         "dtr": dtr_candidates
     }
 
-    description = generate_natural_language_description(tags, api_key)
+    if use_deepseek:
+        description = generate_natural_language_description(tags, api_key=deepseek_key, base_url="https://api.deepseek.com", model="deepseek-chat")
+    else:
+        description = generate_natural_language_description(tags, api_key=api_key)
 
     tags_list = [item for sublist in tags.values() for item in (sublist if isinstance(sublist, list) else [sublist])]  # Flatten
     unique_tags = list(dict.fromkeys(tags_list))
@@ -138,6 +142,14 @@ def gradio_interface():
             type="password"
         )
 
+        deepseek_key_input = gr.Textbox(
+            label="Enter your DeepSeek API Key (Optional)",
+            placeholder="sk-...",
+            type="password"
+        )
+
+        use_deepseek = gr.Checkbox(label="Use DeepSeek API")
+
         with gr.Row():
             action_file = gr.File(label="Upload Action File (Optional)", file_types=[".txt"])
             style_file = gr.File(label="Upload Style File (Optional)", file_types=[".txt"])
@@ -176,7 +188,8 @@ def gradio_interface():
             generate_prompt,
             inputs=[
                 action_file, style_file, artist_files, character_files, dtr_enabled, api_key_input, selected_categories,
-                expression_count, item_count, detail_count, scene_count, angle_count, quality_count, action_count, style_count
+                expression_count, item_count, detail_count, scene_count, angle_count, quality_count, action_count, style_count,
+                use_deepseek, deepseek_key_input
             ],
             outputs=[tags_output, description_output, combined_output],
         )