Spaces:

FineToon
/

Ai-Website-builder

Sleeping

App Files Files Community

AiCoderv2 commited on Sep 14, 2025

Commit

2712381

verified ·

1 Parent(s): b760790

Update app.py

Browse files

Files changed (1) hide show

app.py +95 -45

app.py CHANGED Viewed

@@ -1,14 +1,19 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 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(
@@ -18,63 +23,107 @@ try:
         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")
@@ -95,7 +144,8 @@ with gr.Blocks(theme=gr.themes.Soft()) as app:
                 code_editor = gr.Code(
                     label="HTML Code Editor",
                     language="html",
-                    lines=30
                 )
         with gr.Row():

 import gradio as gr
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForCausalLM
 import torch
 import warnings
 # Suppress warnings
 warnings.filterwarnings("ignore")
+# Try to load models with proper fallbacks
 model_loaded = False
+tokenizer = None
+model = None
+# Try ERNIE model first
 try:
+    print("Loading ERNIE model...")
     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("ERNIE model loaded successfully")
 except Exception as e:
+    print(f"ERNIE model failed: {e}")
+    # Try FLAN-T5 model (seq2seq)
     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("FLAN-T5 model loaded successfully")
     except Exception as e2:
+        print(f"FLAN-T5 model also failed: {e2}")
+        # Final fallback - use a simple template system
+        print("Using template-based generation")
+        model_loaded = False
 def generate_code(prompt):
     """Generate HTML code"""
+    if not model_loaded or not tokenizer or not model:
+        # Template-based fallback
+        return f"""<!DOCTYPE html>
+<html>
+<head>
+    <title>{prompt or 'AI Generated Website'}</title>
+    <style>
+        body {{ font-family: Arial, sans-serif; margin: 40px; background: #f0f0f0; }}
+        .container {{ max-width: 800px; margin: 0 auto; background: white; padding: 30px; border-radius: 10px; }}
+        h1 {{ color: #333; }}
+        p {{ color: #666; }}
+    </style>
+</head>
+<body>
+    <div class="container">
+        <h1>{prompt or 'Generated Website'}</h1>
+        <p>This website was generated based on your description.</p>
+        <button onclick="alert('Hello!')">Click Me</button>
+    </div>
+</body>
+</html>"""
+    try:
+        if "t5" in str(type(model)).lower():
+            # FLAN-T5 handling (seq2seq)
+            full_prompt = f"Create a complete HTML file with CSS and JavaScript for: {prompt}. Return only valid HTML code."
+            inputs = tokenizer(full_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)
+            return result
+        else:
+            # Causal LM handling (ERNIE)
+            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", truncation=True, max_length=512).to("cuda")
+            outputs = model.generate(**inputs, max_new_tokens=1000, temperature=0.7, do_sample=True)
+            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
+    except Exception as e:
+        # Fallback template on error
+        return f"""<!DOCTYPE html>
+<html>
+<head>
+    <title>Error - {prompt or 'AI Generated Website'}</title>
+</head>
+<body>
+    <h1>Generation Error</h1>
+    <p>{str(e)}</p>
+    <p>Using template instead...</p>
+</body>
+</html>"""
 def improve_code(description, current_code):
     """Improve existing code"""
+    if not model_loaded or not tokenizer or not model:
         return current_code
+    try:
+        if "t5" in str(type(model)).lower():
+            # FLAN-T5 handling
+            prompt = f"Improve this HTML code based on the request: {description}\n\nCurrent code:\n{current_code}"
+            inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512)
+            outputs = model.generate(**inputs, max_new_tokens=600, temperature=0.7)
+            return tokenizer.decode(outputs[0], skip_special_tokens=True)
+        else:
+            # ERNIE handling
+            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)
+            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
+    except Exception as e:
+        return current_code
 with gr.Blocks(theme=gr.themes.Soft()) as app:
     gr.Markdown("# AI Website Builder")
                 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():