Spaces:

FineToon
/

Ai-Website-builder

Sleeping

App Files Files Community

AiCoderv2 commited on Sep 14, 2025

Commit

c7b1731

verified ·

1 Parent(s): 32b45b9

Update app.py

Browse files

Files changed (1) hide show

app.py +36 -93

app.py CHANGED Viewed

@@ -1,5 +1,5 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import torch
 import time
 import warnings
@@ -7,81 +7,36 @@ import warnings
 # Suppress warnings
 warnings.filterwarnings("ignore")
-# Load model
 try:
-    print("Loading FLAN-T5 model...")
-    model_name = "google/flan-t5-base"
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
     model_loaded = True
-    print("Model loaded successfully")
 except Exception as e:
-    print(f"Model loading failed: {e}")
     model_loaded = False
 def generate_code_stream(prompt):
     """Stream HTML code generation token by token"""
     if not model_loaded:
-        # Template-based fallback with streaming
-        template = f"""<!DOCTYPE html>
-<html lang="en">
-<head>
-    <meta charset="UTF-8">
-    <meta name="viewport" content="width=device-width, initial-scale=1.0">
-    <title>{prompt or 'AI Generated Website'}</title>
-    <style>
-        * {{ margin: 0; padding: 0; box-sizing: border-box; }}
-        body {{ font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif; background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); min-height: 100vh; padding: 20px; }}
-        .container {{ max-width: 1200px; margin: 0 auto; }}
-        header {{ background: rgba(255,255,255,0.95); padding: 2rem; border-radius: 15px; text-align: center; box-shadow: 0 10px 30px rgba(0,0,0,0.2); }}
-        h1 {{ color: #2c3e50; margin-bottom: 1rem; }}
-        .content {{ background: rgba(255,255,255,0.95); margin: 20px 0; padding: 2rem; border-radius: 15px; box-shadow: 0 10px 30px rgba(0,0,0,0.2); }}
-        button {{ background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); color: white; border: none; padding: 12px 25px; border-radius: 8px; cursor: pointer; font-size: 16px; transition: transform 0.2s; }}
-        button:hover {{ transform: translateY(-2px); }}
-        .features {{ display: grid; grid-template-columns: repeat(auto-fit, minmax(300px, 1fr)); gap: 20px; margin: 2rem 0; }}
-        .feature-card {{ background: #f8f9fa; padding: 1.5rem; border-radius: 10px; border-left: 5px solid #667eea; }}
-    </style>
-</head>
-<body>
-    <div class="container">
-        <header>
-            <h1>{prompt or 'AI Generated Website'}</h1>
-            <p>Created with artificial intelligence</p>
-        </header>
-        <div class="content">
-            <h2>Welcome to Your Generated Site</h2>
-            <p>This website was created based on your description: "{prompt}"</p>
-            <div class="features">
-                <div class="feature-card">
-                    <h3>Feature 1</h3>
-                    <p>Modern design with responsive layout</p>
-                </div>
-                <div class="feature-card">
-                    <h3>Feature 2</h3>
-                    <p>Interactive elements and animations</p>
-                </div>
-                <div class="feature-card">
-                    <h3>Feature 3</h3>
-                    <p>Clean and professional appearance</p>
-                </div>
-            </div>
-            <button onclick="alert('Hello! Thanks for visiting!')">Click Me</button>
-        </div>
-    </div>
-</body>
-</html>"""
-        # Stream character by character
-        for i in range(len(template)):
-            time.sleep(0.005)  # Fast streaming
-            yield template[:i+1]
         return
     try:
-        full_prompt = f"Create a complete modern HTML file with embedded CSS for: {prompt}. Include responsive design, modern styling, and interactive elements. Return only valid HTML code."
-        inputs = tokenizer(full_prompt, return_tensors="pt", truncation=True, max_length=512)
         # Generate with streaming
         outputs = model.generate(
@@ -93,23 +48,19 @@ def generate_code_stream(prompt):
         # Decode and stream character by character
         decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        # Extract HTML if there's extra text
-        if '<!DOCTYPE html>' in decoded:
-            start = decoded.find('<!DOCTYPE html>')
-            decoded = decoded[start:]
         # Stream character by character for smooth effect
-        for i in range(len(decoded)):
             time.sleep(0.003)  # Very fast streaming
-            yield decoded[:i+1]
     except Exception as e:
-        # Fallback streaming
-        error_template = f"<!-- Generation Error: {str(e)} -->\n<!DOCTYPE html>\n<html>\n<head>\n    <title>Error - {prompt or 'AI Generated Website'}</title>\n</head>\n<body>\n    <h1>Generation Failed</h1>\n    <p>Please try again with a different description</p>\n</body>\n</html>"
-        for i in range(len(error_template)):
-            time.sleep(0.005)
-            yield error_template[:i+1]
 def run_code(html_code):
     """Run the generated code in preview"""
@@ -118,28 +69,21 @@ def run_code(html_code):
 def improve_code(description, current_code):
     """Improve existing code"""
     if not model_loaded:
-        return current_code
     try:
         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", truncation=True, max_length=512)
-        outputs = model.generate(**inputs, max_new_tokens=800, temperature=0.7)
-        result = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        # Extract HTML if there's extra text
-        if '<!DOCTYPE html>' in result:
-            start = result.find('<!DOCTYPE html>')
-            result = 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("Generate websites with AI - Streaming Edition")
-    # State to manage current view
-    current_view = gr.State(value="editor")  # "editor" or "preview"
     with gr.Tab("Builder"):
         with gr.Row():
@@ -157,8 +101,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as app:
                 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    <p>Enter a description and click Generate Website</p>\n</body>\n</html>"
                 )
         # Preview area (initially hidden)

 import gradio as gr
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 import time
 import warnings
 # Suppress warnings
 warnings.filterwarnings("ignore")
+# Load the ERNIE model
 try:
+    print("Loading ERNIE Thinking model...")
+    model_name = "baidu/ERNIE-4.5-21B-A3B-Thinking"
+    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name,
+        torch_dtype=torch.bfloat16,
+        device_map="auto",
+        trust_remote_code=True
+    )
     model_loaded = True
+    print("ERNIE Thinking model loaded successfully")
 except Exception as e:
+    print(f"ERNIE model loading failed: {e}")
     model_loaded = False
 def generate_code_stream(prompt):
     """Stream HTML code generation token by token"""
     if not model_loaded:
+        # Error message
+        error_msg = "<!-- Model not loaded -->\n<h1>Model Loading Failed</h1>\n<p>Please check the console for details</p>"
+        for i in range(len(error_msg)):
+            time.sleep(0.01)
+            yield error_msg[:i+1]
         return
     try:
+        full_prompt = f"Create a complete HTML file with embedded CSS and JavaScript for: {prompt}. Return only valid HTML code."
+        inputs = tokenizer(full_prompt, return_tensors="pt", truncation=True, max_length=512).to("cuda")
         # Generate with streaming
         outputs = model.generate(
         # Decode and stream character by character
         decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        result = decoded[len(full_prompt):]
         # Stream character by character for smooth effect
+        for i in range(len(result)):
             time.sleep(0.003)  # Very fast streaming
+            yield result[:i+1]
     except Exception as e:
+        # Error streaming
+        error_msg = f"<!-- Generation Error: {str(e)} -->\n<h1>Generation Failed</h1>\n<p>Please try again</p>"
+        for i in range(len(error_msg)):
+            time.sleep(0.01)
+            yield error_msg[:i+1]
 def run_code(html_code):
     """Run the generated code in preview"""
 def improve_code(description, current_code):
     """Improve existing code"""
     if not model_loaded:
+        return "<!-- Model not loaded -->\n<h1>Model Loading Failed</h1>"
     try:
         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", truncation=True, max_length=512).to("cuda")
+        outputs = model.generate(**inputs, max_new_tokens=800, temperature=0.7, do_sample=True)
+        decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        result = decoded[len(prompt):]
         return result
     except Exception as e:
+        return f"<!-- Error: {str(e)} -->\n{current_code}"
 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")
     with gr.Tab("Builder"):
         with gr.Row():
                 code_editor = gr.Code(
                     label="HTML Code Editor",
                     language="html",
+                    lines=30
                 )
         # Preview area (initially hidden)