Update app.py

app.py

@@ -1,25 +1,20 @@
 import os
 import gradio as gr
 import torch
-from diffusers import StableDiffusionPipeline, DDPMScheduler, Trainer, TrainingArguments, DiffusionPipeline
+from diffusers import StableDiffusionPipeline, DreamBoothTrainer, TrainingArguments, Trainer
+from transformers import CLIPTextModel, CLIPTokenizer
 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
+# تحميل نموذج Stable Diffusion الأساسي
 MODEL_NAME = "runwayml/stable-diffusion-v1-5"
 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 إذا كان موجودًا
@@ -31,81 +26,84 @@ else:
 
 saved_models = {}
 
-# دالة لحفظ بيانات المستخدم
-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 "يجب تسجيل الدخول لحفظ النموذج."
     
     user_model_id = f"user_model_{user_email}"
-    
-    # إعداد مجلد لحفظ الصور
-    user_image_folder = Path(f"user_images/{user_email}")
-    user_image_folder.mkdir(parents=True, exist_ok=True)
+    saved_models[user_email] = user_model_id
 
-    # حفظ الصور في المجلد
-    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))
+    # تأكد من أن الصور مرفوعة
+    if len(images) == 0:
+        return "⚠️ يجب رفع صور للتدريب."
 
-    # تجهيز المعلمات لتدريب DreamBooth
-    accelerator = Accelerator()
+    # تجهيز البيانات الخاصة بالتدريب
+    dataset = load_dataset('imagefolder', data_dir=images)
+
+    # إعداد الـ Trainer للتدريب
+    training_args = TrainingArguments(
+        output_dir="./output",            # مكان حفظ النموذج المدرب
+        per_device_train_batch_size=1,    # حجم الدفعة
+        num_train_epochs=3,               # عدد العصور (epochs)
+        logging_dir='./logs',             # مكان تسجيل السجلات
+        logging_steps=10,
+    )
 
-    # إعداد DreamBooth باستخدام مكتبة diffusers
-    model = StableDiffusionPipeline.from_pretrained(MODEL_NAME)
-    tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-base-patch32")
+    # تحميل Tokenizer و Model من CLIP لتوليد النصوص
+    tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-base-patch16")
+    text_encoder = CLIPTextModel.from_pretrained("openai/clip-vit-base-patch16")
 
-    # تفعيل التدريب باستخدام DreamBooth
+    # إعداد الـ Trainer باستخدام 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"
-        ),
+        model=pipe.unet, 
+        args=training_args, 
         train_dataset=dataset["train"],
-        tokenizer=tokenizer
+        tokenizer=tokenizer,
+        text_encoder=text_encoder,
     )
-    
-    # بداية عملية التدريب
+
+    # بدء عملية التدريب
     trainer.train()
-    
-    # حفظ النموذج المدرب
-    model.save_pretrained(f"user_models/{user_email}")
-    saved_models[user_email] = user_model_id
 
     return f"✅ تم حفظ النموذج بنجاح: {user_model_id}"
 
-# وظيفة إنشاء الصور باستخدام البرومبتات
+# وظيفة لإنشاء الصور باستخدام البرومبتات
 def generate_image(prompt, user_email=None):
     if user_email and user_email in saved_models:
         model_id = saved_models[user_email]
-        model = StableDiffusionPipeline.from_pretrained(f"user_models/{user_email}")  # تحميل النموذج المدرب للمستخدم
-        model.to(device)
-        result = model(prompt).images[0]
+        result = pipe(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("# 🖼️ إنشاء صور مخصصة باستخدام الذكاء الاصطناعي")
     
     with gr.Row():
         with gr.Column():
-            image_input = gr.Files(label="📤 رفع صورك للتدريب")  # ✅ تعديل هنا
+            image_input = gr.Files(label="📤 رفع صورك للتدريب")  # رفع الصور
             user_email_input = gr.Textbox(label="📧 بريدك الإلكتروني (اختياري لحفظ النموذج)")
             train_button = gr.Button("🔧 تدريب النموذج")
             train_output = gr.Textbox(label="🔔 نتيجة التدريب")