Update main.py

main.py

@@ -2,222 +2,178 @@ import sys
 import os
 import types
 import logging
+import subprocess
 from pathlib import Path
 
-import kagglehub
 import torch
 from fastapi import FastAPI
 from fastapi.responses import HTMLResponse
 
-# ── إعداد الـ Logging ─────────────────────────────────────
-logging.basicConfig(
-    level=logging.INFO,
-    format="%(asctime)s | %(levelname)s | %(message)s",
-)
+logging.basicConfig(level=logging.INFO, format="%(asctime)s | %(levelname)s | %(message)s")
 logger = logging.getLogger(__name__)
 
-# ── 0. تجهيز Kaggle Credentials (ضروري على HF Spaces) ─────
-# أضف KAGGLE_USERNAME و KAGGLE_KEY كـ Secrets في إعدادات الـ Space
+# ── 0. تحديث kagglehub (مهم جداً!) ────────────────────────
+try:
+    subprocess.check_call([
+        sys.executable, "-m", "pip", "install", "-q", "--upgrade", "kagglehub>=1.0.2"
+    ])
+    logger.info("✅ kagglehub upgraded")
+except Exception as e:
+    logger.warning(f"⚠️ upgrade failed: {e}")
+
+import kagglehub
+
+# ── 1. Kaggle Credentials ─────────────────────────────────
 kaggle_username = os.environ.get("KAGGLE_USERNAME")
 kaggle_key = os.environ.get("KAGGLE_KEY")
 
 if kaggle_username and kaggle_key:
     os.environ["KAGGLE_USERNAME"] = kaggle_username
     os.environ["KAGGLE_KEY"] = kaggle_key
-    logger.info("✅ Kaggle credentials configured from environment")
+    # اكتب kaggle.json يدوياً (بعض الإصدارات بتقرأ من الملف)
+    kaggle_dir = Path.home() / ".kaggle"
+    kaggle_dir.mkdir(parents=True, exist_ok=True)
+    kaggle_json = kaggle_dir / "kaggle.json"
+    kaggle_json.write_text(f'{{"username":"{kaggle_username}","key":"{kaggle_key}"}}')
+    os.chmod(kaggle_json, 0o600)
+    logger.info("✅ Kaggle credentials configured")
 else:
-    logger.warning("⚠️ KAGGLE_USERNAME/KAGGLE_KEY not found. Public download may fail on HF Spaces.")
+    logger.warning("⚠️ KAGGLE_USERNAME/KAGGLE_KEY not set")
 
-# ── 1. تحميل النموذج من Kaggle ────────────────────────────
-MODEL_HANDLE = "ruicompany/onyx-code-figma/keras/default"
+# ── 2. تحميل النموذج (محاولات متعددة) ─────────────────────
 model_path = None
 
+# محاولة 1: transformers framework (الصحيح للـ .safetensors)
 try:
-    logger.info(f"⬇️ جاري تحميل النموذج: {MODEL_HANDLE}")
-    model_path = kagglehub.model_download(MODEL_HANDLE, force_download=False)
-    logger.info(f"✅ النموذج محمل في: {model_path}")
+    logger.info("⬇️ [Attempt 1] transformers/default")
+    model_path = kagglehub.model_download("ruicompany/onyx-code-figma/transformers/default")
+    logger.info(f"✅ Success: {model_path}")
 except Exception as e:
-    logger.error(f"❌ فشل تحميل النموذج من Kaggle: {e}")
-    
-    # Fallback: جرب مسارات محلية شائعة
-    fallback_paths = [
+    logger.error(f"❌ Attempt 1 failed: {e}")
+
+# محاولة 2: keras/default (القديم - إذا ما غيّرته)
+if not model_path:
+    try:
+        logger.info("⬇️ [Attempt 2] keras/default")
+        model_path = kagglehub.model_download("ruicompany/onyx-code-figma/keras/default")
+        logger.info(f"✅ Success: {model_path}")
+    except Exception as e:
+        logger.error(f"❌ Attempt 2 failed: {e}")
+
+# محاولة 3: dataset_download (أكتر استقراراً)
+if not model_path:
+    try:
+        logger.info("⬇️ [Attempt 3] dataset_download")
+        model_path = kagglehub.dataset_download("ruicompany/onyx-code-figma")
+        logger.info(f"✅ Success: {model_path}")
+    except Exception as e:
+        logger.error(f"❌ Attempt 3 failed: {e}")
+
+# محاولة 4: مسار محلي (إذا ضمّنت المودل بالـ Docker)
+if not model_path:
+    local_paths = [
+        "/app/model",
+        "./model", 
         "/kaggle/working/onyx-code-figma-bundle",
         "./onyx-code-figma-bundle",
-        "./model",
-        "/app/model",
-        "/tmp/onyx-code-figma-bundle",
     ]
-    for p in fallback_paths:
+    for p in local_paths:
         if os.path.exists(p) and any(Path(p).rglob("*.safetensors")):
             model_path = p
-            logger.info(f"✅ استخدام مسار محلي: {model_path}")
+            logger.info(f"✅ [Fallback] Local path: {p}")
             break
 
-# ── 2. فحص محتويات المجلد ─────────────────────────────────
-def inspect_model_dir(path: str):
-    """يفحص الملفات الموجودة في مجلد النموذج"""
-    if not path or not os.path.exists(path):
-        return []
-    items = []
-    for root, dirs, files in os.walk(path):
-        level = root.replace(path, '').count(os.sep)
-        indent = ' ' * 2 * level
-        items.append(f"{indent}{os.path.basename(root)}/")
-        subindent = ' ' * 2 * (level + 1)
-        for file in files:
-            items.append(f"{subindent}{file}")
-    return items
+if not model_path:
+    logger.error("🚨 All download attempts failed!")
+    model_path = None
 
+# ── 3. فحص الملفات ────────────────────────────────────────
 if model_path:
-    logger.info("📁 محتويات مجلد النموذج:")
-    for line in inspect_model_dir(model_path):
-        logger.info(line)
+    logger.info("📁 Files in model path:")
+    for f in sorted(Path(model_path).rglob("*"))[:30]:
+        logger.info(f"   {f.relative_to(model_path)}")
 
-# ── 3. البحث عن onyx_code.py داخل المجلد ──────────────────
+# ── 4. البحث عن onyx_code.py ──────────────────────────────
 onyx_code_dir = None
 if model_path:
     for py_file in Path(model_path).rglob("onyx_code.py"):
         onyx_code_dir = py_file.parent
-        logger.info(f"🔍 Found onyx_code.py at: {py_file}")
+        logger.info(f"🔍 Found onyx_code.py: {py_file}")
         break
 
-# ── 4. فخ IPython (Monkey-patch) محسّن ───────────────────
+# ── 5. Monkey-patch IPython ────────────────────────────────
 captured_html = []
 
 class MockHTML:
-    def __init__(self, data=None, url=None, filename=None, **kwargs):
+    def __init__(self, data=None, **kwargs):
         self.data = data or ""
-        for k, v in kwargs.items():
-            setattr(self, k, v)
 
 def mock_display(*objs, **kwargs):
-    """يلتقط أي كائن HTML يُعرض عبر IPython"""
     for obj in objs:
         if hasattr(obj, 'data') and isinstance(obj.data, str):
             captured_html.append(obj.data)
-            logger.info(f"📄 تم التقاط HTML: {len(obj.data)} حرف")
         elif isinstance(obj, str):
             captured_html.append(obj)
-            logger.info(f"📄 تم التقاط نص: {len(obj)} حرف")
-
-def mock_display_html(*objs, **kwargs):
-    """بديل لـ display_html"""
-    mock_display(*objs, **kwargs)
 
-# إنشاء modules وهمية
 ipython_pkg = types.ModuleType("IPython")
 ipython_display = types.ModuleType("IPython.display")
 ipython_display.HTML = MockHTML
 ipython_display.display = mock_display
-ipython_display.display_html = mock_display_html
+ipython_display.display_html = mock_display
 ipython_display.clear_output = lambda *a, **k: None
 ipython_display.Javascript = lambda *a, **k: MockHTML()
-ipython_display.JSON = lambda *a, **k: MockHTML()
 
 ipython_pkg.display = ipython_display
 sys.modules["IPython"] = ipython_pkg
 sys.modules["IPython.display"] = ipython_display
 
-# ── 5. استيراد onyx_code والتقاط HTML ─────────────────────
-html_content = "<h2>⚠️ لم يتم التقاط HTML</h2>"
-onyx_code_module = None
+# ── 6. استيراد onyx_code ──────────────────────────────────
+html_content = "<h2>⚠️ Onyx Designer not loaded</h2>"
 
 if onyx_code_dir:
     sys.path.insert(0, str(onyx_code_dir))
     try:
         import onyx_code
-        onyx_code_module = onyx_code
-        logger.info("✅ تم استيراد onyx_code بنجاح")
-        
+        logger.info("✅ onyx_code imported")
         if hasattr(onyx_code, 'launch_designer'):
-            logger.info("🎨 تشغيل launch_designer()...")
             onyx_code.launch_designer()
-            if captured_html:
-                html_content = captured_html[-1]
-                logger.info(f"📄 HTML captured: {len(html_content)} chars")
-            else:
-                logger.warning("⚠️ launch_designer() اشتغلت بس ما التقطتش HTML")
-        else:
-            logger.warning("⚠️ onyx_code.launch_designer() غير موجودة")
-            public_funcs = [f for f in dir(onyx_code) if not f.startswith('_')]
-            logger.info(f"🔍 الدوال المتاحة: {public_funcs}")
-            
-    except ImportError as e:
-        logger.error(f"❌ فشل استيراد onyx_code: {e}")
+            html_content = captured_html[-1] if captured_html else html_content
+            logger.info(f"📄 HTML: {len(html_content)} chars")
     except Exception as e:
-        logger.error(f"❌ خطأ أثناء تشغيل launch_designer(): {e}")
+        logger.error(f"❌ Error: {e}")
         import traceback
         logger.error(traceback.format_exc())
-elif model_path:
-    # جرب الاستيراد من المسار الرئيسي مباشرة
-    sys.path.insert(0, model_path)
-    try:
-        import onyx_code
-        onyx_code_module = onyx_code
-        logger.info("✅ تم استيراد onyx_code من model_path")
-        if hasattr(onyx_code, 'launch_designer'):
-            onyx_code.launch_designer()
-            html_content = captured_html[-1] if captured_html else html_content
-    except Exception as e:
-        logger.error(f"❌ فشل استيراد onyx_code من model_path: {e}")
 
-# ── 6. محاولة تحميل النموذج (Transformers) ────────────────
+# ── 7. تحميل Transformers Model ───────────────────────────
 model = None
 tokenizer = None
 
-if model_path:
-    has_pytorch = any(
-        f.suffix in ['.bin', '.safetensors'] or f.name == 'pytorch_model.bin'
-        for f in Path(model_path).rglob('*')
-    )
-    has_keras = any(
-        f.suffix in ['.keras', '.h5'] or f.name == 'saved_model.pb'
-        for f in Path(model_path).rglob('*')
-    )
-    
-    logger.info(f"🔍 PyTorch weights found: {has_pytorch}")
-    logger.info(f"🔍 Keras weights found: {has_keras}")
-
-    if has_pytorch:
-        try:
-            from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
-            
-            quantization_config = BitsAndBytesConfig(
-                load_in_8bit=True,
-                llm_int8_enable_fp32_cpu_offload=True,
-            )
-            
-            model = AutoModelForCausalLM.from_pretrained(
-                model_path,
-                quantization_config=quantization_config,
-                device_map="auto",
-                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
-            )
-            logger.info("✅ تم تحميل النموذج بـ PyTorch/Transformers")
-        except Exception as e:
-            logger.error(f"❌ فشل تحميل النموذج بـ Transformers: {e}")
-            model = None
-
-    elif has_keras:
-        logger.warning("⚠️ النموذج بصيغة Keras. AutoModelForCausalLM لا يدعم Keras مباشرة.")
-        logger.info("💡 استخدم tensorflow.keras.models.load_model() بدلاً من ذلك.")
-        try:
-            import tensorflow as tf
-            model = tf.keras.models.load_model(model_path)
-            logger.info("✅ تم تحميل النموذج بـ TensorFlow/Keras")
-        except Exception as e:
-            logger.error(f"❌ فشل تحميل Keras model: {e}")
+if model_path and any(Path(model_path).rglob("*.safetensors")):
+    try:
+        from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+        
+        quantization_config = BitsAndBytesConfig(
+            load_in_8bit=True,
+            llm_int8_enable_fp32_cpu_offload=True,
+        )
+        
+        model = AutoModelForCausalLM.from_pretrained(
+            model_path,
+            quantization_config=quantization_config,
+            device_map="auto",
+            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)
+        logger.info("✅ Model loaded")
+    except Exception as e:
+        logger.error(f"❌ Model load failed: {e}")
 
-# ── 7. FastAPI ─────────────────────────────────────────────
+# ── 8. FastAPI ────────────────────────────────────────────
 app = FastAPI(title="Onyx Code Figma API")
 
 @app.get("/", response_class=HTMLResponse)
@@ -229,31 +185,9 @@ def health_check():
     return {
         "model_loaded": model is not None,
         "model_path": model_path,
-        "html_captured": len(captured_html) > 0,
         "html_length": len(html_content),
-        "onyx_code_loaded": onyx_code_module is not None,
-    }
-
-@app.get("/model-info")
-def model_info():
-    """يعرض معلومات عن النموذج المحمل"""
-    if not model_path:
-        return {"error": "Model not downloaded"}
-    
-    files = []
-    for f in Path(model_path).rglob('*'):
-        if f.is_file():
-            files.append(str(f.relative_to(model_path)))
-    
-    return {
-        "model_path": model_path,
-        "has_pytorch": has_pytorch if 'has_pytorch' in dir() else None,
-        "has_keras": has_keras if 'has_keras' in dir() else None,
-        "python_files": [str(f.relative_to(model_path)) for f in Path(model_path).rglob("*.py")][:20],
-        "all_files": files[:50],
     }
 
 if __name__ == "__main__":
     import uvicorn
-    logger.info("🚀 تشغيل FastAPI على http://0.0.0.0:7860")
     uvicorn.run(app, host="0.0.0.0", port=7860)