Update app.py

app.py

@@ -1,14 +1,27 @@
 import gradio as gr
 import graphviz
 import os
-from huggingface_hub import InferenceClient
+from transformers import T5Tokenizer, T5ForConditionalGeneration
 from PIL import Image, ImageDraw, ImageFont
+import torch
 
-# --- Constants and Configuration ---
-# NOW USING a Text-Generation model which is more economical with free credits.
-MODEL_ID = "google/flan-t5-large"
+# --- 1. MODEL LOADING (LOCALLY INSIDE THE SPACE) ---
+# No more Inference API! We are loading the model directly.
+print("--- Initializing Local Model ---")
+MODEL_ID = "google/flan-t5-base" # A powerful model small enough to run on a free CPU Space
 
-# The system prompt needs to be a single, direct instruction for this type of model.
+# Check for device
+DEVICE = "cuda:0" if torch.cuda.is_available() else "cpu"
+print(f"--- Using device: {DEVICE} ---")
+
+# Load the tokenizer and model from the Hub into the Space's memory
+tokenizer = T5Tokenizer.from_pretrained(MODEL_ID)
+model = T5ForConditionalGeneration.from_pretrained(MODEL_ID).to(DEVICE)
+print(f"--- Model {MODEL_ID} Initialized Successfully ---")
+
+
+# --- 2. SETUP ---
+# The prompt template for our instruction-tuned 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`).
@@ -20,67 +33,46 @@ Text: "{user_prompt}"
 
 DOT Language Code:"""
 
-# --- Helper Functions ---
+# --- Helper function for placeholder images
 def create_placeholder_image(text="Flowchart will be generated here", size=(600, 800), path="placeholder.png"):
     # (This function remains unchanged)
     try:
         img = Image.new('RGB', size, color=(255, 255, 255))
         draw = ImageDraw.Draw(img)
-        try:
-            font = ImageFont.truetype("DejaVuSans.ttf", 24)
-        except IOError:
-            font = ImageFont.load_default()
-            
+        try: font = ImageFont.truetype("DejaVuSans.ttf", 24)
+        except IOError: font = ImageFont.load_default()
         bbox = draw.textbbox((0, 0), text, font=font)
-        text_width = bbox[2] - bbox[0]
-        text_height = bbox[3] - bbox[1]
+        text_width, text_height = bbox[2] - bbox[0], bbox[3] - bbox[1]
         position = ((size[0] - text_width) / 2, (size[1] - text_height) / 2)
         draw.text(position, text, fill=(200, 200, 200), font=font)
         img.save(path)
         return path
-    except Exception:
-        return None
+    except Exception: return None
 
-# --- Core AI and Rendering Logic ---
-def generate_flowchart(prompt: str, hf_token: str):
+
+# --- 3. CORE AI AND RENDERING LOGIC ---
+def generate_flowchart(prompt: str):
     """
-    Calls the Hugging Face Inference API to generate DOT code and then renders it.
-    Returns the file path of the generated PNG image.
+    Generates a flowchart using the LOCALLY loaded model. No API token is needed.
     """
-    if not hf_token:
-        return create_placeholder_image("Error: Hugging Face API Token is not set.\nPlease add it to your Space's secrets."), None
-
     if not prompt:
         return create_placeholder_image("Please enter a prompt to generate a flowchart."), None
 
     try:
-        # 1. Prepare the full prompt for the text-generation model
+        # 1. Prepare the full prompt and tokenize it
         full_prompt = SYSTEM_PROMPT_TEMPLATE.format(user_prompt=prompt)
+        inputs = tokenizer(full_prompt, return_tensors="pt").input_ids.to(DEVICE)
         
-        client = InferenceClient(model=MODEL_ID, token=hf_token)
-        
-        # --- 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,
-        )
-        # --- 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()
-        if dot_code.startswith("```"):
-            dot_code = dot_code[len("```"):].strip()
-        if dot_code.endswith("```"):
-            dot_code = dot_code[:-len("```")].strip()
-
-        # 2. Render the DOT code using Graphviz
+        # 2. Generate the output from the local model
+        outputs = model.generate(inputs, max_new_tokens=1024, temperature=0.8, do_sample=True)
+        dot_code = tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
+
+        # 3. Clean up the generated code
+        if dot_code.startswith("```dot"): dot_code = dot_code[len("```dot"):].strip()
+        if dot_code.startswith("```"): dot_code = dot_code[len("```"):].strip()
+        if dot_code.endswith("```"): dot_code = dot_code[:-len("```")].strip()
+
+        # 4. Render the DOT code using Graphviz
         graph = graphviz.Source(dot_code)
         output_path = graph.render(os.path.join("outputs", "flowchart"), format='png', cleanup=True)
         
@@ -88,11 +80,12 @@ def generate_flowchart(prompt: str, hf_token: str):
 
     except Exception as e:
         print(f"An error occurred: {e}")
-        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)}"
+        error_message = f"An error occurred during generation.\nThe AI might have produced invalid flowchart code, or another issue occurred.\n\nDetails: {str(e)}"
         return create_placeholder_image(error_message), gr.update(visible=False)
 
-# --- Gradio UI ---
-# (The entire Gradio UI block remains unchanged)
+
+# --- 4. GRADIO UI ---
+# (The Gradio UI block remains mostly unchanged, just removing the token logic)
 css = """
 footer {display: none !important}
 .gradio-container {background-color: #f8f9fa}
@@ -100,13 +93,9 @@ footer {display: none !important}
 """
 
 with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
-    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 "
-        "online AI flowchart generator from text or an intuitive flowchart maker AI, this tool delivers accurate and "
-        "visually engaging results. Discover how AI can enhance your workflow with the best flow chart generator "
-        "AI solution available online."
+        "Our AI Flowchart Generator allows you to create detailed flowcharts instantly. This version runs a self-contained model directly in this Space."
     )
     with gr.Group():
         with gr.Row(equal_height=False):
@@ -121,21 +110,17 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
                 status_display = gr.Markdown("", elem_id="status_display")
             with gr.Column(scale=1):
                 output_image = gr.Image(
-                    label="Generated Flowchart",
-                    type="filepath",
-                    interactive=False,
-                    value=create_placeholder_image(),
-                    height=600,
-                    show_label=False
+                    label="Generated Flowchart", type="filepath", interactive=False,
+                    value=create_placeholder_image(), height=600, show_label=False
                 )
                 download_btn = gr.DownloadButton(
-                    "⬇️ Download",
-                    variant="primary",
-                    visible=False,
+                    "⬇️ Download", variant="primary", visible=False,
                 )
-    def on_generate_click(prompt):
-        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)
+                
+    def on_generate_click(prompt, progress=gr.Progress(track_tqdm=True)):
+        yield (gr.update(interactive=False), gr.update(visible=False), create_placeholder_image("🧠 Thinking... Please wait."), "Generating...")
+        # Note: The 'hf_token' is no longer passed here
+        img_path, download_btn_update = generate_flowchart(prompt)
         yield (gr.update(interactive=True), download_btn_update, img_path, "")
         
     generate_btn.click(
@@ -145,6 +130,5 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
     )
 
 if __name__ == "__main__":
-    if not os.path.exists("outputs"):
-        os.makedirs("outputs")
+    if not os.path.exists("outputs"): os.makedirs("outputs")
     demo.launch()