Spaces:

FineToon
/

Ai-Website-builder

Sleeping

App Files Files Community

AiCoderv2 commited on Sep 14, 2025

Commit

9d6a6b8

verified ·

1 Parent(s): 2712381

Update app.py

Browse files

Files changed (1) hide show

app.py +53 -81

app.py CHANGED Viewed

@@ -1,6 +1,7 @@
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForCausalLM
 import torch
 import warnings
 # Suppress warnings
@@ -11,42 +12,23 @@ 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(
-        model_name,
-        torch_dtype=torch.bfloat16,
-        device_map="auto",
-        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>
@@ -65,39 +47,42 @@ def generate_code(prompt):
     </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"""
@@ -105,29 +90,16 @@ def improve_code(description, current_code):
         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")
-    gr.Markdown("Powered by local AI models")
     with gr.Tab("Builder"):
         with gr.Row():
@@ -154,9 +126,9 @@ with gr.Blocks(theme=gr.themes.Soft()) as app:
         with gr.Row():
             preview = gr.HTML(label="Website Preview")
-    # Event handling
     gen_btn.click(
-        fn=generate_code,
         inputs=desc_input,
         outputs=code_editor
     )

 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForCausalLM
 import torch
+import time
 import warnings
 # Suppress warnings
 tokenizer = None
 model = None
+# Try FLAN-T5 model (more reliable)
 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 e:
+    print(f"FLAN-T5 model failed: {e}")
+    model_loaded = False
+def stream_code(prompt):
+    """Stream HTML code generation token by token"""
     if not model_loaded or not tokenizer or not model:
+        # Template-based fallback with streaming effect
+        template = f"""<!DOCTYPE html>
 <html>
 <head>
     <title>{prompt or 'AI Generated Website'}</title>
     </div>
 </body>
 </html>"""
+        # Stream character by character
+        for i in range(len(template)):
+            time.sleep(0.01)  # Fast streaming
+            yield template[:i+1]
+        return
     try:
+        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)
+        # Generate with streaming
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=88000,
+            temperature=0.7,
+            do_sample=True,
+            output_scores=True,
+            return_dict_in_generate=True
+        )
+        # Decode and stream token by token
+        generated_tokens = outputs.sequences[0]
+        decoded = tokenizer.decode(generated_tokens, skip_special_tokens=True)
+        # Stream character by character for smooth effect
+        for i in range(len(decoded)):
+            time.sleep(0.005)  # Very fast streaming
+            yield decoded[:i+1]
     except Exception as e:
+        # Fallback streaming
+        error_template = f"<!-- Error occurred: {str(e)} -->\n<!DOCTYPE html>\n<html>\n<head>\n    <title>Error</title>\n</head>\n<body>\n    <h1>Generation Failed</h1>\n</body>\n</html>"
+        for i in range(len(error_template)):
+            time.sleep(0.01)
+            yield error_template[:i+1]
 def improve_code(description, current_code):
     """Improve existing code"""
         return current_code
     try:
+        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)
     except Exception as e:
         return current_code
 with gr.Blocks(theme=gr.themes.Soft()) as app:
     gr.Markdown("# AI Website Builder")
+    gr.Markdown("Powered by local AI models - Streaming Generation")
     with gr.Tab("Builder"):
         with gr.Row():
         with gr.Row():
             preview = gr.HTML(label="Website Preview")
+    # Event handling with streaming
     gen_btn.click(
+        fn=stream_code,
         inputs=desc_input,
         outputs=code_editor
     )