Update main.py

main.py

@@ -1,86 +1,37 @@
-import os
-import asyncio
-import json
-import uvicorn
-from fastapi import FastAPI, Request
-from fastapi.responses import StreamingResponse, HTMLResponse
-from gradio_client import Client
+import sys
+import torch
+import kagglehub
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 
-app = FastAPI(title="ONYX Engine")
+model_path = kagglehub.model_download("ruicompany/onyx-code-figma/keras/default")
 
-# تهيئة عميل Gradio للاتصال بـ qwen-onnx-chat مع التوكن الصحيح
-client = Client("ONYX-APP/qwen-onnx-chat", token=os.getenv("HF_TOKEN"))
+# إعداد CPU Offload مع 8-bit
+quantization_config = BitsAndBytesConfig(
+    load_in_8bit=True,
+    llm_int8_enable_fp32_cpu_offload=True,  # ⚡ السر: تفريغ على CPU
+)
 
-@app.get("/")
-async def root():
-    return {"status": "ONYX Engine is running", "model": "google/gemma-4-E2B-it"}
+# تحديد كم الذاكرة لكل جهاز
+max_memory = {
+    0: "0GB",           # لا تستخدم GPU نهائياً
+    "cpu": "6GB",       # خصص 6GB من RAM للنموذج
+}
 
-@app.post("/predict")
-async def generate(request: Request):
-    data = await request.json()
-    messages = data.get("messages", [])
-    
-    # 1. استخراج الـ system_prompt الأساسي، وتجميع باقي الرسائل
-    system_prompt = ""
-    filtered_chat = []
-    for msg in messages:
-        role = msg.get("role")
-        content = msg.get("content", "")
-        if role == "system":
-            if system_prompt:
-                system_prompt += "\n" + content
-            else:
-                system_prompt = content
-        else:
-            filtered_chat.append(msg)
-                
-    # 2. تحديد رسالة المستخدم الأخيرة وتاريخ المحادثة وقبلهما
-    user_message = ""
-    history_msgs = []
-    if filtered_chat:
-        user_message = filtered_chat[-1].get("content", "")
-        history_msgs = filtered_chat[:-1]
-                    
-    # 3. بناء الـ history بشكل أزواج متتالية [ [user, assistant], [user, assistant], ... ]
-    history = []
-    i = 0
-    while i < len(history_msgs):
-        if history_msgs[i].get("role") == "user":
-            u_content = history_msgs[i].get("content", "")
-            a_content = ""
-            if i + 1 < len(history_msgs) and history_msgs[i+1].get("role") == "assistant":
-                a_content = history_msgs[i+1].get("content", "")
-                i += 1
-            history.append([u_content, a_content])
-        i += 1
+model = AutoModelForCausalLM.from_pretrained(
+    model_path,
+    quantization_config=quantization_config,
+    device_map="auto",
+    max_memory=max_memory,
+    offload_buffers=True,
+    offload_folder="offload",       # مجلد التفريغ المؤقت
+    offload_state_dict=True,
+    torch_dtype=torch.float32,      # CPU يحتاج float32
+    trust_remote_code=True,
+)
+tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
 
-    async def generate_chunks():
-        try:
-            def run_client():
-                # [تعديل] تمرير المعاملات بالترتيب (Positional) حصراً لأن السبيس لا يدعم الـ keyword arguments
-                return client.predict(
-                    user_message,
-                    history,
-                    system_prompt
-                )
-                                    
-            response_text = await asyncio.to_thread(run_client)
-                                    
-            # إرسال النص الناتج على دفعات متتالية لتوافق الـ StreamingResponse
-            for char in str(response_text):
-                yield char
-                await asyncio.sleep(0.01)
-                                
-        except Exception as e:
-            yield f" بعتذر منك بس فيك تشغل التطبيق بكرة لانو المودل الاساسي قيد التطوير "
+print("✅ النموذج يعمل على CPU + 8-bit Quantization")
 
-    headers = {
-        "Content-Type": "text/event-stream",
-        "Cache-Control": "no-cache",
-        "Connection": "keep-alive",
-        "X-Accel-Buffering": "no"
-    }
-    return StreamingResponse(generate_chunks(), headers=headers)
-
-if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+sys.path.insert(0, model_path)
+from onyx_code import launch_designer
+launch_designer()