Update app.py

app.py

@@ -1,22 +1,18 @@
-import gradio as gr
+# app.py
 from transformers import pipeline
+import gradio as gr
 
-# Load the text-to-text generation model
-generator = pipeline("text2text-generation", model="google/flan-t5-small")
+# Use a better model trained for HTML generation
+generator = pipeline("text2text-generation", model="mrm8488/t5-base-finetuned-web-code")
 
-# Define the function to generate website code
-def generate_website(prompt):
-    # Instruction prompt to ensure web-style code output
-    full_prompt = f"Generate a modern HTML, CSS, and JavaScript website. {prompt}"
-    
-    result = generator(full_prompt, max_length=1024, do_sample=True)[0]['generated_text']
+def generate_site(prompt):
+    result = generator(prompt, max_length=2048, do_sample=True)[0]["generated_text"]
     return result
 
-# Build a simple Gradio UI
-gr.Interface(
-    fn=generate_website,
-    inputs=gr.Textbox(label="Describe the website you want", placeholder="e.g. A modern cafe website using beige and brown colors..."),
-    outputs=gr.Textbox(label="Generated Website Code"),
-    title="AI Website Code Generator",
-    description="Enter a prompt to generate a full website (HTML/CSS/JS)."
-).launch()
+iface = gr.Interface(fn=generate_site,
+                     inputs=gr.Textbox(lines=10, placeholder="Enter your website prompt..."),
+                     outputs="text",
+                     title="AI Website Generator",
+                     description="Enter a prompt to generate a modern HTML/CSS/JS website")
+
+iface.launch()