Update app.py

app.py

@@ -1,24 +1,53 @@
+import google.generativeai as genai
+import gradio as gr
 from html2image import Html2Image
 import uuid
 import os
-import google.generativeai as genai
-import gradio as gr
 
+# -----------------------
 # Configure Gemini API
+# -----------------------
 genai.configure(api_key=os.environ.get("GEMINI_API_KEY"))
 model = genai.GenerativeModel("gemini-2.0-flash")
 
-hti = Html2Image()
+# -----------------------
+# Html2Image Setup
+# -----------------------
+hti = Html2Image(executable_path="/usr/bin/chromium")
 
+# -----------------------
+# Function: Generate Mockup Image
+# -----------------------
 def generate_mockup(user_prompt):
     system_prompt = (
         "You are a Power Apps mockup generator. "
         "Convert user descriptions into complete HTML + CSS code that looks like a Power Apps mobile screen. "
         "Do NOT include JSON or explanations. Only return valid HTML + CSS."
     )
-    response = model.generate_content(system_prompt + "\n" + user_prompt)
-    return response.text.strip()
 
+    try:
+        # Generate HTML
+        response = model.generate_content(system_prompt + "\n" + user_prompt)
+        html_code = response.text.strip()
+
+        # Unique file names
+        unique_id = str(uuid.uuid4())
+        html_file = f"mockup_{unique_id}.html"
+        img_file = f"mockup_{unique_id}.png"
+
+        with open(html_file, "w", encoding="utf-8") as f:
+            f.write(html_code)
+
+        # Convert HTML -> PNG
+        hti.screenshot(html_file=html_file, save_as=img_file)
+
+        return img_file
+    except Exception as e:
+        return f"Error: {e}"
+
+# -----------------------
+# Gradio UI
+# -----------------------
 with gr.Blocks() as demo:
     gr.Markdown("## ⚡ Power Apps Mockup Generator (Gemini + Gradio)")
     user_input = gr.Textbox(
@@ -27,8 +56,8 @@ with gr.Blocks() as demo:
         placeholder="Example: A login screen with username, password, and login button"
     )
     generate_btn = gr.Button("Generate Mockup")
-    html_output = gr.HTML(label="Mockup Preview")
+    image_output = gr.Image(label="Mockup Preview")
 
-    generate_btn.click(fn=generate_mockup, inputs=user_input, outputs=html_output)
+    generate_btn.click(fn=generate_mockup, inputs=user_input, outputs=image_output)
 
 demo.launch(server_name="0.0.0.0", server_port=7860)