Update app.py

app.py

@@ -1,10 +1,16 @@
 import os
 import gradio as gr
 import torch
-from diffusers import StableDiffusionPipeline
+from diffusers import StableDiffusionPipeline, DDPMScheduler, Trainer, TrainingArguments, DiffusionPipeline
+from datasets import load_dataset
+from transformers import CLIPTextModel, CLIPTokenizer, PreTrainedModel
 from werkzeug.security import generate_password_hash, check_password_hash
 import json
 import time
+from pathlib import Path
+from PIL import Image
+from diffusers import DreamBoothTrainer, DreamBoothPipeline
+from accelerate import Accelerator
 
 # تحميل نموذج Stable Diffusion
 MODEL_NAME = "runwayml/stable-diffusion-v1-5"
@@ -13,7 +19,7 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 pipe = StableDiffusionPipeline.from_pretrained(MODEL_NAME)
 pipe.to(device)
 
-# قاعدة بيانات للمستخدمين والنماذج المحفوظة
+# قاعدة بيانات للمستخدمين والنماذج المحفوظة (سيتم حفظها في ملف JSON)
 user_data_file = "user_data.json"
 
 # تحميل البيانات من ملف JSON إذا كان موجودًا
@@ -25,16 +31,59 @@ else:
 
 saved_models = {}
 
-# وظيفة تدريب النموذج (تمت إزالة شريط التقدم)
-def train_model(user_email, images):
+# دالة لحفظ بيانات المستخدم
+def save_users_data():
+    with open(user_data_file, "w") as f:
+        json.dump(users, f)
+
+# وظيفة تدريب النموذج على الصور المرفوعة باستخدام DreamBooth أو LoRA
+def train_model(user_email, images, progress):
     if not user_email:
-        return "❌ يجب تسجيل الدخول لحفظ النموذج."
+        return "يجب تسجيل الدخول لحفظ النموذج."
     
     user_model_id = f"user_model_{user_email}"
-    saved_models[user_email] = user_model_id
+    
+    # إعداد مجلد لحفظ الصور
+    user_image_folder = Path(f"user_images/{user_email}")
+    user_image_folder.mkdir(parents=True, exist_ok=True)
+
+    # حفظ الصور في المجلد
+    for img in images:
+        img.save(user_image_folder / img.name)
+    
+    # إعداد بيانات التدريب (يمكنك تخصيص هذه الطريقة لتناسب النموذج الخاص بك)
+    image_paths = [str(user_image_folder / img.name) for img in images]
+    dataset = load_dataset('imagefolder', data_dir=str(user_image_folder))
+
+    # تجهيز المعلمات لتدريب DreamBooth
+    accelerator = Accelerator()
 
-    # محاكاة عملية التدريب
-    time.sleep(5)  # تأخير زمني لمحاكاة التدريب (يمكن تعديله)
+    # إعداد DreamBooth باستخدام مكتبة diffusers
+    model = StableDiffusionPipeline.from_pretrained(MODEL_NAME)
+    tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-base-patch32")
+
+    # تفعيل التدريب باستخدام DreamBooth
+    trainer = DreamBoothTrainer(
+        model=model,
+        args=TrainingArguments(
+            output_dir=f"user_models/{user_email}",
+            per_device_train_batch_size=4,
+            num_train_epochs=1,
+            gradient_accumulation_steps=2,
+            logging_dir='./logs',
+            logging_steps=10,
+            report_to="none"
+        ),
+        train_dataset=dataset["train"],
+        tokenizer=tokenizer
+    )
+    
+    # بداية عملية التدريب
+    trainer.train()
+    
+    # حفظ النموذج المدرب
+    model.save_pretrained(f"user_models/{user_email}")
+    saved_models[user_email] = user_model_id
 
     return f"✅ تم حفظ النموذج بنجاح: {user_model_id}"
 
@@ -42,31 +91,14 @@ def train_model(user_email, images):
 def generate_image(prompt, user_email=None):
     if user_email and user_email in saved_models:
         model_id = saved_models[user_email]
-        result = pipe(prompt).images[0]
+        model = StableDiffusionPipeline.from_pretrained(f"user_models/{user_email}")  # تحميل النموذج المدرب للمستخدم
+        model.to(device)
+        result = model(prompt).images[0]
     else:
         result = pipe(prompt).images[0]
 
     return result
 
-# وظيفة تسجيل المستخدمين
-def register(email, password):
-    if email in users:
-        return "⚠️ البريد الإلكتروني مسجل بالفعل!"
-    users[email] = {"password": generate_password_hash(password)}
-    save_users_data()
-    return "✅ تم التسجيل بنجاح!"
-
-# وظيفة تسجيل الدخول
-def login(email, password):
-    if email not in users or not check_password_hash(users[email]["password"], password):
-        return "❌ البريد الإلكتروني أو كلمة المرور غير صحيحة!"
-    return f"✅ تم تسجيل الدخول: {email}"
-
-# حفظ بيانات المستخدمين إلى ملف JSON
-def save_users_data():
-    with open(user_data_file, "w") as f:
-        json.dump(users, f)
-
 # واجهة Gradio
 with gr.Blocks() as demo:
     gr.Markdown("# 🖼️ إنشاء صور مخصصة باستخدام الذكاء الاصطناعي")
@@ -77,13 +109,14 @@ with gr.Blocks() as demo:
             user_email_input = gr.Textbox(label="📧 بريدك الإلكتروني (اختياري لحفظ النموذج)")
             train_button = gr.Button("🔧 تدريب النموذج")
             train_output = gr.Textbox(label="🔔 نتيجة التدريب")
+            progress_bar = gr.Progress()  # بدون label هنا
  
         with gr.Column():
             prompt_input = gr.Textbox(label="✏️ أدخل البرومبت لإنشاء صورة")
             generate_button = gr.Button("🎨 إنشاء صورة")
             output_image = gr.Image(label="📷 الصورة الناتجة")
  
-    train_button.click(train_model, inputs=[user_email_input, image_input], outputs=train_output)
+    train_button.click(train_model, inputs=[user_email_input, image_input, progress_bar], outputs=train_output)
     generate_button.click(generate_image, inputs=[prompt_input, user_email_input], outputs=output_image)
  
     gr.Markdown("### 🔑 تسجيل الدخول / التسجيل")