Update app.py

app.py

@@ -5,23 +5,24 @@ from huggingface_hub import InferenceClient
 from PIL import Image, ImageDraw, ImageFont
 
 # --- Constants and Configuration ---
-# Using a powerful open-source instruction-tuned model from Mistral AI
+# NOW USING a Text-Generation model which is more economical with free credits.
 MODEL_ID = "google/flan-t5-large"
 
-# System prompt to instruct the AI model
-# This is the most crucial part for getting reliable, structured output.
-SYSTEM_PROMPT = """You are an expert at creating flowcharts. A user will provide a text prompt, and your task is to generate the flowchart description in the Graphviz DOT language.
-
+# The system prompt needs to be a single, direct instruction for this type of model.
+SYSTEM_PROMPT_TEMPLATE = """Task: Generate a flowchart description in the Graphviz DOT language based on the following text.
 Your response MUST be ONLY the Graphviz DOT language source code for a directed graph (digraph).
 - The graph should be top-to-bottom (`rankdir=TB`).
 - Use rounded boxes for process steps (`shape=box, style="rounded,filled", fillcolor="#EAEAFB"`).
 - Use diamonds for decision points (`shape=diamond, fillcolor="#F9EED5"`).
 - Use ellipses for start and end nodes (`shape=ellipse, fillcolor="#D5D6F9"`).
-- Do not include any explanations, comments, or markdown formatting like ```dot ... ```. The entire response must be valid DOT code."""
+
+Text: "{user_prompt}"
+
+DOT Language Code:"""
 
 # --- Helper Functions ---
 def create_placeholder_image(text="Flowchart will be generated here", size=(600, 800), path="placeholder.png"):
-    """Creates a placeholder or error image with text."""
+    # (This function remains unchanged)
     try:
         img = Image.new('RGB', size, color=(255, 255, 255))
         draw = ImageDraw.Draw(img)
@@ -53,18 +54,24 @@ def generate_flowchart(prompt: str, hf_token: str):
         return create_placeholder_image("Please enter a prompt to generate a flowchart."), None
 
     try:
-        # 1. Call the LLM to get the DOT language code
+        # 1. Prepare the full prompt for the text-generation model
+        full_prompt = SYSTEM_PROMPT_TEMPLATE.format(user_prompt=prompt)
+        
         client = InferenceClient(model=MODEL_ID, token=hf_token)
-        response = client.chat_completion(
-            messages=[
-                {"role": "system", "content": SYSTEM_PROMPT},
-                {"role": "user", "content": prompt},
-            ],
-            max_tokens=1500,
+        
+        # --- THIS IS THE KEY CHANGE ---
+        # Use client.text_generation instead of client.chat_completion
+        dot_code = client.text_generation(
+            prompt=full_prompt,
+            max_new_tokens=1024, # Flan-T5 can generate a lot, let's give it space
             temperature=0.7,
+            do_sample=True,
         )
-        dot_code = response.choices[0].message.content.strip()
+        # --- END OF KEY CHANGE ---
         
+        # The result is a direct string, so we just need to strip it.
+        dot_code = dot_code.strip()
+
         # Sometimes the model still adds markdown, let's strip it just in case
         if dot_code.startswith("```dot"):
             dot_code = dot_code[len("```dot"):].strip()
@@ -84,8 +91,8 @@ def generate_flowchart(prompt: str, hf_token: str):
         error_message = f"An error occurred.\nThis could be due to an invalid API token,\nan issue with the AI model, or invalid generated DOT code.\n\nDetails: {str(e)}"
         return create_placeholder_image(error_message), gr.update(visible=False)
 
-
 # --- Gradio UI ---
+# (The entire Gradio UI block remains unchanged)
 css = """
 footer {display: none !important}
 .gradio-container {background-color: #f8f9fa}
@@ -93,9 +100,7 @@ footer {display: none !important}
 """
 
 with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
-    # Get the Hugging Face token from environment variables
     hf_token = os.environ.get("HF_TOKEN")
-
     gr.Markdown("# AI Flowchart Generator")
     gr.Markdown(
         "Our AI Flowchart Generator allows you to create detailed flowcharts instantly. Whether you need a free "
@@ -103,7 +108,6 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
         "visually engaging results. Discover how AI can enhance your workflow with the best flow chart generator "
         "AI solution available online."
     )
-
     with gr.Group():
         with gr.Row(equal_height=False):
             with gr.Column(scale=1):
@@ -114,9 +118,7 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
                 )
                 with gr.Row():
                     generate_btn = gr.Button("✨ Generate", variant="primary")
-                
                 status_display = gr.Markdown("", elem_id="status_display")
-
             with gr.Column(scale=1):
                 output_image = gr.Image(
                     label="Generated Flowchart",
@@ -131,25 +133,10 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
                     variant="primary",
                     visible=False,
                 )
-
-    # --- Event Listeners ---
     def on_generate_click(prompt):
-        # Provide user feedback that generation is in progress
-        yield (
-            gr.update(interactive=False), # Disable button
-            gr.update(visible=False),     # Hide download button
-            create_placeholder_image("🧠 Generating with AI... Please wait."),
-            "Generating... this can take up to 30 seconds."
-        )
-        # Call the main generation function
+        yield (gr.update(interactive=False), gr.update(visible=False), create_placeholder_image("🧠 Generating with AI... Please wait."), "Generating... this can take up to 30 seconds.")
         img_path, download_btn_update = generate_flowchart(prompt, hf_token)
-        # Update UI with the result
-        yield (
-            gr.update(interactive=True), # Re-enable button
-            download_btn_update,         # Show/update download button
-            img_path,
-            "" # Clear status message
-        )
+        yield (gr.update(interactive=True), download_btn_update, img_path, "")
         
     generate_btn.click(
         fn=on_generate_click,