Update app.py

app.py

@@ -1,13 +1,14 @@
-# Updated UI Generator with Better Open-Source Model
+# Updated UI Generator with Better Open-Source Model and DartPad/Expo Preview Support
 
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextGenerationPipeline
 import torch
+import re
 
-# Use a better model for code generation, such as DeepSeek-Coder or Codestral
-model_id = "deepseek-ai/deepseek-coder-6.7b-instruct"  # Change to another if desired
+# Use a smaller model for CPU support
+model_id = "deepseek-ai/deepseek-coder-1.3b-instruct"
 
-# Ensure torch uses CPU if no GPU
+# Use CPU or GPU if available
 device = 0 if torch.cuda.is_available() else -1
 
 # Load tokenizer and model
@@ -15,19 +16,29 @@ print("Loading model...")
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32)
 
-# Build the text generation pipeline
+# Build the generation pipeline
 generator = TextGenerationPipeline(model=model, tokenizer=tokenizer, device=device)
 
 def generate_ui(platform, framework, ui_prompt):
-    prompt = f"""
-You are an expert mobile app developer.
-Generate a complete {framework} UI code snippet for a {platform} app based on the description:
-"{ui_prompt}"
-Include all required imports, a main method, and best practices for UI structure.
+    prompt_template = f"""
+You are a professional mobile developer.
+Create a full {framework} UI code for {platform} based on this request:
 """
-    response = generator(prompt, max_new_tokens=512, do_sample=True, temperature=0.7)[0]['generated_text']
-    # Trim the echoed prompt and just return the generated code
-    return response.split(""""""")[-1].strip()
+    full_prompt = prompt_template + f"\n{ui_prompt}\n" + "\nEnsure to include imports, widgets, and app structure."
+    response = generator(full_prompt, max_new_tokens=512, do_sample=True, temperature=0.7)[0]['generated_text']
+
+    # Extract only code block (removes echoed prompt if needed)
+    code_match = re.search(r"(?s)(import .*?)$", response)
+    code = code_match.group(1).strip() if code_match else response.strip()
+    return code, generate_preview_link(code, framework)
+
+def generate_preview_link(code, framework):
+    if framework == "Flutter":
+        return f"https://dartpad.dev/?id=custom&code={gr.utils.sanitize_html(code)[:1500]}"
+    elif framework == "React Native":
+        return f"https://snack.expo.dev/?code={gr.utils.sanitize_html(code)[:1500]}"
+    else:
+        return "Preview not available for this framework."
 
 interface = gr.Interface(
     fn=generate_ui,
@@ -36,9 +47,12 @@ interface = gr.Interface(
         gr.Dropdown(["Flutter", "Kotlin XML", "SwiftUI", "React Native"], label="Framework"),
         gr.Textbox(lines=4, label="UI Prompt", placeholder="e.g. Login screen with email & password, dark theme")
     ],
-    outputs=gr.Code(label="Generated UI Code"),
+    outputs=[
+        gr.Code(label="Generated UI Code"),
+        gr.Textbox(label="Preview Link")
+    ],
     title="Prompt-to-UI Code Generator",
-    description="Generate Android/iOS UI code in Flutter, SwiftUI, XML, or React Native by just describing the layout."
+    description="Describe your screen and get UI code + preview for Flutter/React Native apps."
 )
 
 interface.launch()