Update app.py

app.py

@@ -1,33 +1,36 @@
 import gradio as gr
-import os, zipfile, tempfile
+import os, zipfile, tempfile, json
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
-import json
 
 # ---------------------------
-# 1️⃣ Load Local Model
+# 1️⃣ Load Public Code Model
 # ---------------------------
-model_name = "TheBloke/WizardCoder-7B-GPTQ"
+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)
+generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
 
 # ---------------------------
 # 2️⃣ Website Generation Function
 # ---------------------------
 def generate_website(prompt):
-    # Generate multi-file code as JSON
+    # Instruction for multi-file code in JSON format
     instruction = f"""
-Generate a full website project structure based on the user prompt. 
-Provide the response in JSON format where key = file path, value = file content. 
-Include folders: css/, js/, assets/, db/ if needed. User prompt: {prompt}
+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}
 """
-    gen = generator(instruction, max_new_tokens=2048)[0]['generated_text']
 
-    # Parse JSON safely
+    # 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:
-        files_dict = {"index.html": "<!-- Failed to generate code -->"}
+        # 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()
@@ -57,7 +60,7 @@ iface = gr.Interface(
     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 with folders (css/, js/, assets/, db/). Download as ZIP."
+    description="Single Python file app. Generates multi-file website locally using public HF model. Download as ZIP."
 )
 
 iface.launch()