Update app.py

app.py

@@ -8,37 +8,46 @@ from threading import Thread
 from playwright.async_api import async_playwright
 from playwright_stealth import stealth
 
-# --- OS Setup: Playwright Install ---
+# --- Playwright Installation Check ---
+try:
+    import playwright
+except ImportError:
+    os.system("pip install playwright playwright-stealth")
+    
+# Space startup par browser install karein
 os.system("playwright install chromium")
 
 # --- AI Model Setup ---
-# Hum Llama-3.2-1B use kar rahe hain taaki RAM crash na ho
+# Qwen 2.5 1.5B - Best balance for ZeroGPU RAM
 MODEL_ID = "Qwen/Qwen2.5-1.5B-Instruct"
 
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-# Model ko CPU par load karenge, @spaces.GPU isse khud GPU par le jayega
 model = AutoModelForCausalLM.from_pretrained(
-    MODEL_ID, 
-    torch_dtype=torch.bfloat16, 
-    device_map=None
+    MODEL_ID,
+    torch_dtype=torch.bfloat16, # torch_dtype fix
+    device_map=None # Manual control
 ).to("cpu")
 
 @spaces.GPU
 def chat_fn(message, history):
-    # Convert Gradio history to Llama format
     messages = []
-    for user_msg, assistant_msg in history:
-        messages.append({"role": "user", "content": user_msg})
-        messages.append({"role": "assistant", "content": assistant_msg})
+    for h in history:
+        messages.append({"role": "user", "content": h[0]})
+        messages.append({"role": "assistant", "content": h[1]})
     messages.append({"role": "user", "content": message})
     
-    # Move model to GPU inside the decorated function
-    model.to("cuda")
+    model.to("cuda") # Move to GPU only when needed
+    
+    input_ids = tokenizer.apply_chat_template(
+        messages, 
+        tokenize=True, 
+        add_generation_prompt=True, 
+        return_tensors="pt"
+    ).to("cuda")
     
-    input_ids = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt").to("cuda")
     streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+    gen_kwargs = dict(input_ids=input_ids, streamer=streamer, max_new_tokens=1024, temperature=0.7)
     
-    gen_kwargs = dict(input_ids=input_ids, streamer=streamer, max_new_tokens=512, temperature=0.7)
     thread = Thread(target=model.generate, kwargs=gen_kwargs)
     thread.start()
     
@@ -47,60 +56,72 @@ def chat_fn(message, history):
         response += new_text
         yield response
 
-# --- Browser Logic ---
+# --- Improved Browser Logic ---
 async def browser_task(url, selector=None, text=None, action="nav"):
     async with async_playwright() as p:
-        browser = await p.chromium.launch(headless=True)
-        context = await browser.new_context(viewport={'width': 1024, 'height': 768})
+        # Browser launch with args to prevent sandbox crashes
+        browser = await p.chromium.launch(headless=True, args=["--no-sandbox", "--disable-setuid-sandbox"])
+        context = await browser.new_context(
+            viewport={'width': 1280, 'height': 720},
+            user_agent="Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/122.0.0.0 Safari/537.36"
+        )
         page = await context.new_page()
         await stealth(page)
         
         try:
-            await page.goto(url, wait_until="networkidle", timeout=30000)
+            # Timeout barha diya hai
+            await page.goto(url, wait_until="load", timeout=60000)
+            
             if action == "click" and selector:
+                await page.wait_for_selector(selector, timeout=5000)
                 await page.click(selector)
-                await asyncio.sleep(1)
+                await asyncio.sleep(2)
             elif action == "type" and selector and text:
+                await page.wait_for_selector(selector, timeout=5000)
                 await page.fill(selector, text)
                 await page.keyboard.press("Enter")
-                await asyncio.sleep(2)
+                await asyncio.sleep(3)
             
             screenshot = await page.screenshot()
             await browser.close()
             return screenshot
         except Exception as e:
+            print(f"Browser Error: {e}")
             await browser.close()
             return None
 
 def web_run(url, sel, txt, action):
-    return asyncio.run(browser_task(url, sel, txt, action))
+    try:
+        return asyncio.run(browser_task(url, sel, txt, action))
+    except Exception as e:
+        return None
 
-# --- Gradio UI ---
-with gr.Blocks() as demo:
-    gr.Markdown("## Dual AI: Smart Chat + Stealth Browser")
+# --- Custom UI with Gradio ---
+with gr.Blocks(theme=gr.themes.Default()) as demo:
+    gr.HTML("<h1 style='text-align:center;'>🚀 AI Agent + Web Stealth Browser</h1>")
     
     with gr.Row():
-        # Left side: Chat
-        with gr.Column():
-            chatbot = gr.ChatInterface(fn=chat_fn, type="messages")
+        with gr.Column(scale=1):
+            # Standard Chatbot logic
+            chat_ui = gr.ChatInterface(fn=chat_fn, type="messages")
             
-        # Right side: Browser
-        with gr.Column():
-            url_box = gr.Textbox(label="URL", value="https://www.google.com")
+        with gr.Column(scale=1):
+            url_input = gr.Textbox(label="🌐 Website URL", value="https://google.com")
             with gr.Row():
-                sel_box = gr.Textbox(label="Selector", placeholder="#id or .class")
-                txt_box = gr.Textbox(label="Text to type")
+                sel_input = gr.Textbox(label="🔍 Selector (ID/Class)", placeholder="e.g. input[name='q']")
+                txt_input = gr.Textbox(label="⌨️ Text to Type", placeholder="Search query...")
             
             with gr.Row():
-                btn_go = gr.Button("🌐 Visit")
-                btn_click = gr.Button("🖱️ Click")
-                btn_type = gr.Button("⌨️ Type")
+                btn_visit = gr.Button("Visit", variant="primary")
+                btn_click = gr.Button("Click")
+                btn_type = gr.Button("Type & Enter")
             
-            browser_output = gr.Image(label="Browser View")
+            browser_img = gr.Image(label="Live View", type="numpy")
 
-    # Browser Button Events
-    btn_go.click(fn=lambda u: web_run(u, None, None, "nav"), inputs=url_box, outputs=browser_output)
-    btn_click.click(fn=lambda u, s: web_run(u, s, None, "click"), inputs=[url_box, sel_box], outputs=browser_output)
-    btn_type.click(fn=lambda u, s, t: web_run(u, s, t, "type"), inputs=[url_box, sel_box, txt_box], outputs=browser_output)
+    # Handlers
+    btn_visit.click(lambda u: web_run(u, None, None, "nav"), [url_input], browser_img)
+    btn_click.click(lambda u, s: web_run(u, s, None, "click"), [url_input, sel_input], browser_img)
+    btn_type.click(lambda u, s, t: web_run(u, s, t, "type"), [url_input, sel_input, txt_input], browser_img)
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()