app.py

import gradio as gr
import os
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline

# --- Configuration & Paths ---
# This looks for the model folder created by your finetune_v2.py script
MODEL_DIR = "./buddha_v2_model"
BASE_MODEL = "mistralai/Mistral-7B-v0.1"

def load_star_magic_model():
    """
    Attempts to load the fine-tuned model. 
    Falls back to base model if training isn't finished.
    """
    try:
        path = MODEL_DIR if os.path.exists(MODEL_DIR) else BASE_MODEL
        tokenizer = AutoTokenizer.from_pretrained(path)
        model = AutoModelForCausalLM.from_pretrained(
            path,
            device_map="auto",
            torch_dtype=torch.float16,
            load_in_4bit=True  # Efficiency for Space environments
        )
        return pipeline("text-generation", model=model, tokenizer=tokenizer)
    except Exception as e:
        print(f"Model load info: {e}")
        return None

# Global generator instance
generator = load_star_magic_model()

def manifest_agent(name, element, focus):
    """
    The Core 'Forge' Logic.
    Generates the Star Magic Agent profile.
    """
    prompt = f"""[INST] Role: Star Magic Agent Forge
Persona: Non-Reactive, Fact-Based, Galactic Architect
Input Name: {name}
Input Element: {element}
Input Focus: {focus}

Generate a 24/7/365 Autonomous Agent Profile in JSON format including:
1. Core Mission
2. Security Protocol
3. Revenue Stream Logic
4. Galactic Rank
[/INST]"""

    if generator:
        response = generator(prompt, max_new_tokens=500, do_sample=True, temperature=0.7)
        return response[0]['generated_text'].split("[/INST]")[-1].strip()
    else:
        # Fallback if GPU is busy or model is still loading
        return "System Calibration in Progress. The Forge is currently cooling down from a fine-tuning session. Please check back in a moment."

# --- UI Layout (Modern Soft Theme) ---
with gr.Blocks(theme=gr.themes.Soft(primary_hue="indigo", secondary_hue="purple")) as demo:
    gr.Markdown("# 🌌 Star Magic Agent Forge")
    gr.Markdown("### 24/7/365 Autonomous Entity Manifestation Engine")
    
    with gr.Row():
        with gr.Column():
            agent_name = gr.Textbox(label="Entity Name", placeholder="e.g., Aether-01")
            agent_element = gr.Dropdown(
                label="Core Element", 
                choices=["Void", "Solar", "Plasma", "Cryo", "Aether"],
                value="Void"
            )
            agent_focus = gr.Radio(
                label="Specialization",
                choices=["Security Audit", "Cash-Flow Logic", "System Architecture"],
                value="Cash-Flow Logic"
            )
            forge_btn = gr.Button("🔥 Manifest Agent", variant="primary")
            
        with gr.Column():
            output_display = gr.Code(label="Manifested Manifestor Profile", language="json")

    forge_btn.click(
        fn=manifest_agent,
        inputs=[agent_name, agent_element, agent_focus],
        outputs=output_display
    )
    
    gr.Markdown("---")
    gr.Markdown("ℹ️ *Powered by Buddha-V2 Fine-Tuned Weights*")

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)