import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
import warnings
# Suppress warnings
warnings.filterwarnings("ignore")
# Global variables for model
tokenizer = None
model = None
model_loaded = False
def load_model():
"""Load model on demand"""
global tokenizer, model, model_loaded
if model_loaded:
return True
try:
# Try CodeLlama 7B
print("Loading CodeLlama 7B model...")
model_name = "codellama/CodeLlama-7b-hf"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
model_name,
torch_dtype=torch.float16,
device_map="auto"
)
model_loaded = True
print("Model loaded successfully")
return True
except Exception as e:
print(f"Model loading failed: {e}")
return False
def generate_code_fast(prompt):
"""Instant HTML code generation"""
# Load model if not loaded
if not load_model():
return "\n\n\n Error\n\n\n Model Loading Failed
\n\n"
try:
# Minimal prompt for instant generation
full_prompt = f"HTML for {prompt}:"
inputs = tokenizer(full_prompt, return_tensors="pt", truncation=True, max_length=128).to("cuda" if torch.cuda.is_available() else "cpu")
# Instant generation with minimal processing
outputs = model.generate(
**inputs,
max_new_tokens=500,
temperature=0.7,
do_sample=True
)
# Decode result
decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
result = decoded[len(full_prompt):]
# Ensure valid HTML
if '' not in result:
result = f"\n\n\n {prompt or 'Generated Site'}\n\n\n {prompt or 'Content'}
\n\n"
return result
except Exception as e:
return "\n\n\n Error\n\n\n Generation Failed
\n\n"
def run_code(html_code):
"""Run the generated code in preview"""
return html_code
def improve_code(description, current_code):
"""Improve existing code instantly"""
# Load model if not loaded
if not load_model():
return current_code
try:
# Minimal prompt for instant improvement
prompt = f"Improve: {description}"
inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=128).to("cuda" if torch.cuda.is_available() else "cpu")
outputs = model.generate(**inputs, max_new_tokens=400, temperature=0.7)
decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
result = decoded[len(prompt):]
# Ensure valid HTML
if '' in result:
start = result.find('')
return result[start:]
return result
except Exception as e:
return current_code
with gr.Blocks(theme=gr.themes.Soft()) as app:
gr.Markdown("# AI Website Builder")
gr.Markdown("Instant Code Generation")
with gr.Tab("Builder"):
with gr.Row():
with gr.Column(scale=1):
desc_input = gr.Textbox(
label="Describe your website",
placeholder="e.g., Portfolio site",
lines=2
)
gen_btn = gr.Button("Generate Instantly", variant="primary")
imp_btn = gr.Button("Quick Improve")
run_btn = gr.Button("Run Website")
with gr.Column(scale=2):
code_editor = gr.Code(
label="HTML Code Editor",
language="html",
lines=20
)
# Preview area (initially hidden)
with gr.Row(visible=False) as preview_row:
gr.Markdown("### Live Preview")
with gr.Row(visible=False) as preview_content:
preview = gr.HTML()
# Back to editor button (initially hidden)
with gr.Row(visible=False) as back_row:
back_btn = gr.Button("Back to Editor")
# Event handling - instant generation
gen_btn.click(
fn=generate_code_fast,
inputs=desc_input,
outputs=code_editor
)
imp_btn.click(
fn=improve_code,
inputs=[desc_input, code_editor],
outputs=code_editor
)
run_btn.click(
fn=lambda html: (html, gr.Row(visible=True), gr.Row(visible=True), gr.Row(visible=True)),
inputs=code_editor,
outputs=[preview, preview_row, preview_content, back_row]
).then(
fn=run_code,
inputs=code_editor,
outputs=preview
)
back_btn.click(
fn=lambda: (gr.Row(visible=False), gr.Row(visible=False), gr.Row(visible=False)),
outputs=[preview_row, preview_content, back_row]
)
code_editor.change(
fn=lambda x: x,
inputs=code_editor,
outputs=preview
)
app.launch()