Update main.py

main.py

@@ -1,4 +1,6 @@
 import sys
+import os
+import types
 import kagglehub
 import torch
 from fastapi import FastAPI
@@ -7,76 +9,59 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 
 app = FastAPI()
 
-# 1. تحميل النموذج
-model_path = kagglehub.model_download(
-    "ruicompany/onyx-code-figma/keras/default"
-)
+# ── 1. تحميل النموذج ───────────────────────────────────────
+model_path = kagglehub.model_download("ruicompany/onyx-code-figma/keras/default")
 
 quantization_config = BitsAndBytesConfig(
     load_in_8bit=True,
     llm_int8_enable_fp32_cpu_offload=True,
 )
 
-max_memory = {
-    0: "0GB",
-    "cpu": "6GB",
-}
-
 model = AutoModelForCausalLM.from_pretrained(
     model_path,
     quantization_config=quantization_config,
     device_map="auto",
-    max_memory=max_memory,
+    max_memory={0: "0GB", "cpu": "6GB"},
     offload_buffers=True,
     offload_folder="offload",
     offload_state_dict=True,
     torch_dtype=torch.float32,
     trust_remote_code=True,
 )
-
 tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
-print("✅ النموذج يعمل على CPU + 8-bit Quantization")
-
-# 2. استدعاء onyx_code والتقاط الـ HTML
-sys.path.insert(0, model_path)
-import onyx_code
+print("✅ النموذج جاهز")
 
+# ── 2. فخ IPython (Monkey-patch) ───────────────────────────
 captured_html = []
 
-
-# اعتراض دوال العرض الخاصة بـ IPython داخل onyx_code
 def mock_display(obj, *args, **kwargs):
     if hasattr(obj, "data"):
         captured_html.append(obj.data)
-    elif isinstance(obj, str):
-        captured_html.append(obj)
-
 
 class MockHTML:
-
     def __init__(self, data):
         self.data = data
-        captured_html.append(data)
 
+# ننشئ الـ modules الوهمية قبل أي استيراد
+ipython_pkg = types.ModuleType("IPython")
+ipython_display = types.ModuleType("IPython.display")
+ipython_display.HTML = MockHTML
+ipython_display.display = mock_display
+ipython_pkg.display = ipython_display
 
-onyx_code.display = mock_display
-onyx_code.HTML = MockHTML
+sys.modules["IPython"] = ipython_pkg
+sys.modules["IPython.display"] = ipython_display
 
-# تشغيل الدالة لالتقاط الواجهة
-result = onyx_code.launch_designer()
+# ── 3. استيراد onyx_code والتقاط HTML ─────────────────────
+sys.path.insert(0, model_path)
+import onyx_code
 
-# استخراج النص التفاعلي للـ HTML
-if captured_html:
-    html_content = captured_html[-1]
-elif hasattr(result, "data"):
-    html_content = result.data
-elif result is not None:
-    html_content = str(result)
-else:
-    html_content = "<h2>لم يتم العثور على كود HTML من الدالة</h2>"
+onyx_code.launch_designer()  # الـ mock يلتقط HTML تلقائياً
 
+html_content = captured_html[-1] if captured_html else "<h2>⚠️ لم يتم التقاط HTML</h2>"
+print(f"📄 HTML captured: {len(html_content)} chars")
 
-# 3. عرض الصفحة بـ FastAPI
+# ── 4. FastAPI ─────────────────────────────────────────────
 @app.get("/", response_class=HTMLResponse)
 def read_root():
     return html_content