Update app.py

app.py

@@ -1,27 +1,51 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
 
-tokenizer = AutoTokenizer.from_pretrained("deepseek-ai/deepseek-coder-6.7b-instruct", trust_remote_code=True)
+MODEL_ID = "deepseek-ai/deepseek-coder-1.3b-instruct"
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
 model = AutoModelForCausalLM.from_pretrained(
-    "deepseek-ai/deepseek-coder-6.7b-instruct", 
-    trust_remote_code=True
+    MODEL_ID,
+    trust_remote_code=True  # keep it, but NO device_map to avoid needing accelerate
 )
 
+PROMPT_TEMPLATE = "<|user|>\n{prompt}\n<|assistant|>\n"
 
-generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
+def generate_website_code(user_prompt, max_tokens=512, temperature=0.7, top_p=0.95):
+    prompt = PROMPT_TEMPLATE.format(prompt=user_prompt)
+    inputs = tokenizer(prompt, return_tensors="pt")
+    with torch.no_grad():
+        output_ids = model.generate(
+            **inputs,
+            max_new_tokens=max_tokens,
+            do_sample=True,
+            temperature=temperature,
+            top_p=top_p,
+            pad_token_id=tokenizer.eos_token_id
+        )
+    text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+    code = text.split("<|assistant|>")[-1].strip()
+    return code, code  # (code box, html preview)
 
-def generate_website_code(prompt):
-    formatted_prompt = f"<|user|>\n{prompt}\n<|assistant|>\n"
-    result = generator(formatted_prompt, max_new_tokens=1024, do_sample=True, temperature=0.7)
-    code = result[0]['generated_text'].split("<|assistant|>")[-1].strip()
-    return code
+with gr.Blocks() as demo:
+    gr.Markdown("# Web X — Turn your website ideas into code with AI.")
 
-demo = gr.Interface(
-    fn=generate_website_code,
-    inputs=gr.Textbox(lines=3, label="Describe your website idea"),
-    outputs=gr.Code(label="Generated HTML/CSS/JS Code"),
-    title="WebX — AI Website Generator",
-    description="Turn your website ideas into code using DeepSeek Coder."
-)
+    with gr.Row():
+        prompt = gr.Textbox(
+            label="Describe your website idea",
+            lines=3,
+            placeholder="e.g. Portfolio with navbar, about, projects grid and contact form"
+        )
+        generate_btn = gr.Button("Generate")
+
+    code_out = gr.Code(label="Generated HTML/CSS/JS")
+    preview = gr.HTML(label="Live Preview")
+
+    generate_btn.click(
+        fn=generate_website_code,
+        inputs=prompt,
+        outputs=[code_out, preview]
+    )
 
 demo.launch()