Spaces:

FineToon
/

Ai-Website-builder

Sleeping

App Files Files Community

AiCoderv2 commited on Sep 14, 2025

Commit

b760790

verified ·

1 Parent(s): 893bd1f

Update app.py

Browse files

Files changed (1) hide show

app.py +52 -88

app.py CHANGED Viewed

@@ -4,11 +4,11 @@ import torch
 import warnings
 # Suppress warnings
-warnings.filterwarnings("ignore", category=FutureWarning, module="huggingface_hub")
-warnings.filterwarnings("ignore", category=UserWarning, module="transformers")
 try:
-    # Download and load the model
     model_name = "baidu/ERNIE-4.5-21B-A3B-Thinking"
     tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
     model = AutoModelForCausalLM.from_pretrained(
@@ -18,94 +18,84 @@ try:
         trust_remote_code=True
     )
     model_loaded = True
 except Exception as e:
     print(f"Model loading failed: {e}")
-    model_loaded = False
 def generate_code(prompt):
-    """Generate HTML code using ERNIE model"""
     if not model_loaded:
-        return "<!-- Model not available -->\n<h1>Model Loading Failed</h1>"
-    full_prompt = f"Create a complete single HTML file with embedded CSS and JavaScript for: {prompt}. Return only valid HTML code."
-    inputs = tokenizer(full_prompt, return_tensors="pt").to("cuda")
-    with torch.no_grad():
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=1000,
-            temperature=0.7,
-            top_p=0.9,
-            do_sample=True
-        )
-    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    result = generated_text[len(full_prompt):]
-    # Extract HTML if there's extra text
-    if '<!DOCTYPE html>' in result:
-        start = result.find('<!DOCTYPE html>')
-        return result[start:]
-    return result
 def improve_code(description, current_code):
     """Improve existing code"""
     if not model_loaded:
         return current_code
-    prompt = f"Improve this HTML code based on: {description}\n\nCurrent code:\n{current_code}\n\nReturn only the improved HTML code."
-    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
-    with torch.no_grad():
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=800,
-            temperature=0.7,
-            top_p=0.9,
-            do_sample=True
-        )
-    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    result = generated_text[len(prompt):]
-    # Extract HTML if there's extra text
-    if '<!DOCTYPE html>' in result:
-        start = result.find('<!DOCTYPE html>')
-        return result[start:]
-    return result
 with gr.Blocks(theme=gr.themes.Soft()) as app:
     gr.Markdown("# AI Website Builder")
-    gr.Markdown("Powered by baidu/ERNIE-4.5-21B-A3B-Thinking (Local Model)")
     with gr.Tab("Builder"):
         with gr.Row():
             with gr.Column(scale=1):
-                gr.Markdown("### Instructions")
                 desc_input = gr.Textbox(
                     label="Describe your website",
-                    placeholder="e.g., A modern portfolio with dark theme and animated cards",
                     lines=4
                 )
                 gen_btn = gr.Button("Generate Website", variant="primary")
                 imp_btn = gr.Button("Improve Current Code")
-                gr.Markdown("### Tips")
-                gr.Markdown("""
-                - Be specific about layout and features
-                - Mention color schemes and styling
-                - Describe interactive elements needed
-                - Request responsive design for mobile
-                """)
             with gr.Column(scale=2):
                 code_editor = gr.Code(
                     label="HTML Code Editor",
                     language="html",
-                    lines=30,
-                    value="<!DOCTYPE html>\n<html>\n<head>\n    <title>AI Generated Website</title>\n</head>\n<body>\n    <h1>Your website will appear here</h1>\n</body>\n</html>"
                 )
         with gr.Row():
@@ -114,32 +104,6 @@ with gr.Blocks(theme=gr.themes.Soft()) as app:
         with gr.Row():
             preview = gr.HTML(label="Website Preview")
-    with gr.Tab("About"):
-        gr.Markdown("""
-        # AI Website Builder
-        This tool uses Baidu's ERNIE 4.5 model to generate complete websites from text descriptions.
-        ## Features:
-        - Local model execution (no API keys required)
-        - Real-time code generation
-        - Live website preview
-        - Code editing and improvement
-        - Syntax-highlighted editor
-        ## How to Use:
-        1. Describe your desired website in detail
-        2. Click "Generate Website" to create code
-        3. Edit the generated code if needed
-        4. See live preview of your website
-        5. Use "Improve Current Code" to refine specific aspects
-        ## Model Information:
-        - Model: baidu/ERNIE-4.5-21B-A3B-Thinking
-        - Framework: Hugging Face Transformers
-        - Execution: Local inference
-        """)
     # Event handling
     gen_btn.click(
         fn=generate_code,

 import warnings
 # Suppress warnings
+warnings.filterwarnings("ignore")
+# Try to load model with fallback
+model_loaded = False
 try:
     model_name = "baidu/ERNIE-4.5-21B-A3B-Thinking"
     tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
     model = AutoModelForCausalLM.from_pretrained(
         trust_remote_code=True
     )
     model_loaded = True
+    print("Model loaded successfully")
 except Exception as e:
     print(f"Model loading failed: {e}")
+    # Fallback to a working model
+    try:
+        model_name = "google/flan-t5-base"
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        model = AutoModelForCausalLM.from_pretrained(model_name)
+        model_loaded = True
+        print("Fallback model loaded successfully")
+    except Exception as e2:
+        print(f"Fallback model also failed: {e2}")
 def generate_code(prompt):
+    """Generate HTML code"""
     if not model_loaded:
+        return f"<!-- Model not available -->\n<!DOCTYPE html>\n<html>\n<head>\n    <title>{prompt or 'Generated Site'}</title>\n</head>\n<body>\n    <h1>Model Loading Failed</h1>\n    <p>Using template instead</p>\n</body>\n</html>"
+    if "flan-t5" in str(type(model)):
+        # FLAN-T5 handling
+        full_prompt = f"Create a complete HTML file for: {prompt}. Include CSS and JS."
+        inputs = tokenizer(full_prompt, return_tensors="pt")
+        outputs = model.generate(**inputs, max_new_tokens=800)
+        result = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return result
+    else:
+        # ERNIE handling
+        full_prompt = f"Create a complete single HTML file with embedded CSS and JavaScript for: {prompt}. Return only valid HTML code."
+        inputs = tokenizer(full_prompt, return_tensors="pt").to("cuda")
+        outputs = model.generate(**inputs, max_new_tokens=1000, temperature=0.7)
+        result = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        result = result[len(full_prompt):]
+        if '<!DOCTYPE html>' in result:
+            start = result.find('<!DOCTYPE html>')
+            return result[start:]
+        return result
 def improve_code(description, current_code):
     """Improve existing code"""
     if not model_loaded:
         return current_code
+    if "flan-t5" in str(type(model)):
+        prompt = f"Improve this HTML: {description}\n\n{current_code}"
+        inputs = tokenizer(prompt, return_tensors="pt")
+        outputs = model.generate(**inputs, max_new_tokens=600)
+        return tokenizer.decode(outputs[0], skip_special_tokens=True)
+    else:
+        prompt = f"Improve this HTML code based on: {description}\n\nCurrent code:\n{current_code}\n\nReturn only the improved HTML code."
+        inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
+        outputs = model.generate(**inputs, max_new_tokens=800, temperature=0.7)
+        result = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        result = result[len(prompt):]
+        if '<!DOCTYPE html>' in result:
+            start = result.find('<!DOCTYPE html>')
+            return result[start:]
+        return result
 with gr.Blocks(theme=gr.themes.Soft()) as app:
     gr.Markdown("# AI Website Builder")
+    gr.Markdown("Powered by local AI models")
     with gr.Tab("Builder"):
         with gr.Row():
             with gr.Column(scale=1):
                 desc_input = gr.Textbox(
                     label="Describe your website",
+                    placeholder="e.g., A modern portfolio with dark theme",
                     lines=4
                 )
                 gen_btn = gr.Button("Generate Website", variant="primary")
                 imp_btn = gr.Button("Improve Current Code")
             with gr.Column(scale=2):
                 code_editor = gr.Code(
                     label="HTML Code Editor",
                     language="html",
+                    lines=30
                 )
         with gr.Row():
         with gr.Row():
             preview = gr.HTML(label="Website Preview")
     # Event handling
     gen_btn.click(
         fn=generate_code,