app.py

import gradio as gr
import os, zipfile, tempfile, json
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline

# ---------------------------
# 1️⃣ Load Public Code Model
# ---------------------------
model_name = "Salesforce/codegen-350M-multi"  # public model, no token required
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)
generator = pipeline("text-generation", model=model, tokenizer=tokenizer)

# ---------------------------
# 2️⃣ Website Generation Function
# ---------------------------
def generate_website(prompt):
    # Instruction for multi-file code in JSON format
    instruction = f"""
Generate a full website project based on the user prompt.
Return the result as JSON where key = file path and value = file content.
Include folders: css/, js/, assets/, db/ if needed.
User prompt: {prompt}
"""

    # Generate code text
    gen = generator(instruction, max_new_tokens=1024)[0]['generated_text']

    # Try parsing JSON from generated text
    try:
        files_dict = json.loads(gen)
    except:
        # fallback: single index.html if parsing fails
        files_dict = {"index.html": "<!-- Failed to generate valid code -->"}

    # Create temp folder and save files
    temp_dir = tempfile.mkdtemp()
    for fname, content in files_dict.items():
        file_path = os.path.join(temp_dir, fname)
        os.makedirs(os.path.dirname(file_path), exist_ok=True)
        with open(file_path, "w", encoding="utf-8") as f:
            f.write(content)

    # Zip the folder
    zip_path = os.path.join(temp_dir, "website.zip")
    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zipf:
        for root, _, files in os.walk(temp_dir):
            for file in files:
                if file != os.path.basename(zip_path):
                    file_path = os.path.join(root, file)
                    arcname = os.path.relpath(file_path, temp_dir)
                    zipf.write(file_path, arcname=arcname)

    return zip_path

# ---------------------------
# 3️⃣ Gradio UI
# ---------------------------
iface = gr.Interface(
    fn=generate_website,
    inputs=gr.Textbox(lines=4, placeholder="Describe your website: PHP/HTML/JS, admin panel, MySQL..."),
    outputs=gr.File(label="Download Generated Website ZIP"),
    title="Local AI Multi-file Website Generator",
    description="Single Python file app. Generates multi-file website locally using public HF model. Download as ZIP."
)

iface.launch()