Spaces:

FineToon
/

Ai-Website-builder

Sleeping

App Files Files Community

AiCoderv2 commited on Sep 14

Commit

1febf1f

verified ·

1 Parent(s): f1607ac

Update app.py

Browse files

Files changed (1) hide show

app.py +35 -39

app.py CHANGED Viewed

@@ -19,15 +19,14 @@ def load_model():
         return True
     try:
-        # Try CodeLlama 7B (faster loading)
         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",
-            low_cpu_mem_usage=True
         )
         model_loaded = True
         print("Model loaded successfully")
@@ -37,100 +36,97 @@ def load_model():
         return False
 def generate_code_fast(prompt):
-    """Fast HTML code generation with max tokens"""
     # Load model if not loaded
     if not load_model():
-        return "<!-- Model loading failed -->\n<h1>Model Error</h1>\n<p>Could not load AI model</p>"
     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" if torch.cuda.is_available() else "cpu")
-        # Fast generation with maximum tokens
         outputs = model.generate(
             **inputs,
-            max_new_tokens=88000,  # 88k tokens
             temperature=0.7,
-            do_sample=True,
-            top_p=0.9,
-            repetition_penalty=1.2
         )
         # Decode result
         decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
         result = decoded[len(full_prompt):]
-        # Clean up result
-        if result.startswith('"') and result.endswith('"'):
-            result = result[1:-1]
-        # Ensure it's valid HTML
-        if not result.strip().startswith('<!DOCTYPE html>'):
-            result = f"<!DOCTYPE html>\n<html>\n<head>\n    <title>{prompt or 'Generated Site'}</title>\n</head>\n<body>\n    <h1>{prompt or 'Generated Content'}</h1>\n</body>\n</html>"
         return result
     except Exception as e:
-        return f"<!-- Error: {str(e)} -->\n<h1>Generation Failed</h1>\n<p>Please try again</p>"
 def run_code(html_code):
     """Run the generated code in preview"""
     return html_code
 def improve_code(description, current_code):
-    """Improve existing code"""
     # Load model if not loaded
     if not load_model():
-        return "<!-- Model not loaded -->\n<h1>Model Error</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" if torch.cuda.is_available() else "cpu")
-        outputs = model.generate(**inputs, max_new_tokens=88000, temperature=0.7, do_sample=True)
         decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
         result = decoded[len(prompt):]
-        # Clean up result
-        if result.startswith('"') and result.endswith('"'):
-            result = result[1:-1]
         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("Ultra-Fast Generation - 88K Token Limit")
     with gr.Tab("Builder"):
         with gr.Row():
             with gr.Column(scale=1):
                 desc_input = gr.Textbox(
                     label="Describe your website",
-                    placeholder="e.g., Portfolio with dark theme",
-                    lines=3
                 )
-                gen_btn = gr.Button("Generate Website", variant="primary")
-                imp_btn = gr.Button("Improve Code")
-                run_btn = gr.Button("Run Website", variant="secondary")
             with gr.Column(scale=2):
                 code_editor = gr.Code(
                     label="HTML Code Editor",
                     language="html",
-                    lines=25
                 )
         # 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(label="Website Preview")
         # Back to editor button (initially hidden)
         with gr.Row(visible=False) as back_row:
-            back_btn = gr.Button("Back to Editor", variant="secondary")
-    # Event handling - ultra fast generation
     gen_btn.click(
         fn=generate_code_fast,
         inputs=desc_input,

         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 False
 def generate_code_fast(prompt):
+    """Instant HTML code generation"""
     # Load model if not loaded
     if not load_model():
+        return "<!DOCTYPE html>\n<html>\n<head>\n    <title>Error</title>\n</head>\n<body>\n    <h1>Model Loading Failed</h1>\n</body>\n</html>"
     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 '<!DOCTYPE html>' not in result:
+            result = f"<!DOCTYPE html>\n<html>\n<head>\n    <title>{prompt or 'Generated Site'}</title>\n</head>\n<body>\n    <h1>{prompt or 'Content'}</h1>\n</body>\n</html>"
         return result
     except Exception as e:
+        return "<!DOCTYPE html>\n<html>\n<head>\n    <title>Error</title>\n</head>\n<body>\n    <h1>Generation Failed</h1>\n</body>\n</html>"
 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 '<!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")
+    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,