Update app.py

app.py

@@ -1,18 +1,29 @@
-from huggingface_hub import InferenceClient
-import os
-api_key=os.getenv("HF_TOKEN")
 import gradio as gr
 from huggingface_hub import InferenceClient
 import os
+import json
 
 # Initialize Hugging Face Inference Client
 api_key = os.getenv("HF_TOKEN")
 client = InferenceClient(api_key=api_key)
 
-def chat_with_model(user_input):
+# Load or initialize system prompts
+PROMPTS_FILE = "system_prompts.json"
+if os.path.exists(PROMPTS_FILE):
+    with open(PROMPTS_FILE, "r") as file:
+        system_prompts = json.load(file)
+else:
+    system_prompts = {"default": "You are a good image generation prompt engineer for diffuser image generation models"}
+
+def save_prompts():
+    """Save the current system prompts to a JSON file."""
+    with open(PROMPTS_FILE, "w") as file:
+        json.dump(system_prompts, file, indent=4)
+
+def chat_with_model(user_input, system_prompt):
     """Send user input to the model and return its response."""
     messages = [
-        {"role": "system", "content": "You are a good image generation prompt engineer for diffuser image generation models"},
+        {"role": "system", "content": system_prompt},
         {"role": "user", "content": user_input}
     ]
 
@@ -29,19 +40,38 @@ def chat_with_model(user_input):
     except Exception as e:
         return f"Error: {str(e)}"
 
+def update_prompt(name, content):
+    """Update or add a new system prompt."""
+    system_prompts[name] = content
+    save_prompts()
+    return f"System prompt '{name}' saved."
+
+def get_prompt(name):
+    """Retrieve a system prompt by name."""
+    return system_prompts.get(name, "")
+
 # Gradio Interface
 with gr.Blocks() as demo:
     gr.Markdown("## Hugging Face Chatbot with Gradio")
 
     with gr.Row():
         with gr.Column():
-            user_input = gr.Textbox(label="Enter your prompt", placeholder="Describe the character or request a detailed description...")
+            system_prompt_name = gr.Dropdown(choices=list(system_prompts.keys()), label="Select System Prompt")
+            system_prompt_content = gr.TextArea(label="System Prompt", value="", lines=4)
+            save_prompt_button = gr.Button("Save System Prompt")
+
+            user_input = gr.TextArea(label="Enter your prompt", placeholder="Describe the character or request a detailed description...", lines=4)
             submit_button = gr.Button("Generate")
         
         with gr.Column():
-            output = gr.Textbox(label="Model Response", interactive=False)
+            output = gr.TextArea(label="Model Response", interactive=False, lines=10)
+
+    def load_prompt(name):
+        return get_prompt(name)
 
-    submit_button.click(chat_with_model, inputs=[user_input], outputs=[output])
+    system_prompt_name.change(load_prompt, inputs=[system_prompt_name], outputs=[system_prompt_content])
+    save_prompt_button.click(update_prompt, inputs=[system_prompt_name, system_prompt_content], outputs=[])
+    submit_button.click(chat_with_model, inputs=[user_input, system_prompt_content], outputs=[output])
 
 # Run the app
 demo.launch()