Update app.py

app.py

@@ -1,74 +1,57 @@
-import spaces  # <--- يجب أن يكون هذا هو السطر رقم 1 دائماً وأبداً
+import spaces  # <--- يجب أن يبقى في السطر الأول
 import gradio as gr
-import torch
 import cv2
 import numpy as np
 from PIL import Image
-from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, LCMScheduler
-from controlnet_aux import CannyDetector
 import os
 import shutil
 import tempfile 
 import datetime
-import ffmpeg 
+import ffmpeg
+# استيراد مكتبات ModelScope الخاصة بـ DDColor
+from modelscope.pipelines import pipeline
+from modelscope.utils.constant import Tasks
 
 # ==========================================
-# 1. إعدادات النموذج
+# 1. إعداد نموذج DDColor الاحترافي
 # ==========================================
-print("⏳ Loading Models...")
-
-# تحديد نوع البيانات (ملاحظة: مع ZeroGPU التحديد يتم لاحقاً، لكن نجهزه هنا)
-# ملاحظة: لا تستخدم .to('cuda') هنا خارج الدالة في ZeroGPU
-torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+print("⏳ Loading DDColor Professional Model...")
 
+# تحميل خط الأنابيب (Pipeline) الخاص بالتلوين
+# نحدد device='gpu' ليعمل مع ZeroGPU عند استدعائه
 try:
-    # تحميل ControlNet
-    controlnet_model = ControlNetModel.from_pretrained(
-        "lllyasviel/sd-controlnet-canny", torch_dtype=torch_dtype
-    )
-    
-    # تحميل Stable Diffusion 1.5
-    model_id = "runwayml/stable-diffusion-v1-5"
-    pipe = StableDiffusionControlNetPipeline.from_pretrained(
-        model_id, controlnet=controlnet_model, torch_dtype=torch_dtype
+    ddcolor_pipeline = pipeline(
+        Tasks.image_colorization, 
+        model='damo/cv_ddcolor_image-colorization',
+        device='gpu' 
     )
-    
-    # تفعيل LCM للسرعة
-    print("⚡ Injecting LCM-LoRA...")
-    pipe.load_lora_weights("latent-consistency/lcm-lora-sdv1-5")
-    pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
-    
-    print("✅ Models loaded into RAM (waiting for GPU allocation).")
-    
+    print("✅ DDColor Model loaded successfully.")
 except Exception as e:
-    print(f"❌ Error loading models: {e}")
-    # لن نوقف البرنامج هنا، سنتركه يحاول العمل
-    pass
-
-canny_processor = CannyDetector()
+    print(f"❌ Error loading DDColor model: {e}")
+    ddcolor_pipeline = None
 
 # ==========================================
-# 2. دالة المعالجة (مع تفعيل ZeroGPU)
+# 2. دالة المعالجة (الاحترافية)
 # ==========================================
 
-@spaces.GPU(duration=120)  # نمنح الدالة وقتاً كافياً (120 ثانية)
-def colorize_video_multistyle(video_file, prompt, style_choice, steps=5):
+@spaces.GPU(duration=180) # نمنح وقتاً كافياً للفيديوهات
+def colorize_video_professional(video_file):
     if not video_file:
         return None
-
-    # === نقل النموذج إلى GPU الآن فقط (داخل الدالة) ===
-    print("🚀 Moving models to GPU...")
-    pipe.to("cuda")
     
-    # تحسينات الذاكرة
-    pipe.enable_attention_slicing()
+    if ddcolor_pipeline is None:
+        raise gr.Error("فشل تحميل النموذج. يرجى مراجعة السجلات.")
+
+    print("🚀 Starting professional colorization on ZeroGPU...")
 
     timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
-    output_temp_video_no_audio = os.path.join(tempfile.gettempdir(), f"temp_vis_{timestamp}.mp4")
-    final_output_name = f"colored_lcm_{timestamp}.mp4"
+    temp_frames_dir = os.path.join(tempfile.gettempdir(), f"frames_dd_{timestamp}")
+    os.makedirs(temp_frames_dir, exist_ok=True)
     
-    # --- استخراج الصوت ---
-    audio_path = os.path.join(tempfile.gettempdir(), f"temp_audio_{timestamp}.aac")
+    final_output_name = f"colored_ddcolor_{timestamp}.mp4"
+    audio_path = os.path.join(tempfile.gettempdir(), f"audio_dd_{timestamp}.aac")
+
+    # --- 1. استخراج الصوت (إن وجد) ---
     audio_exists = False
     try:
         (
@@ -81,111 +64,84 @@ def colorize_video_multistyle(video_file, prompt, style_choice, steps=5):
     except ffmpeg.Error:
         print("⚠️ Warning: No audio found or extraction failed.")
 
-    # --- قراءة الفيديو ---
+    # --- 2. قراءة الفيديو ومعالجة الإطارات ---
     cap = cv2.VideoCapture(video_file)
-    if not cap.isOpened():
-        return None
-        
     fps = cap.get(cv2.CAP_PROP_FPS)
-    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-    
-    # تجهيز البرومبت
-    style_prompts = {
-        "Auto Color": "photorealistic, 8k, masterpiece, cinematic lighting",
-        "Vivid": "vibrant colors, high saturation, pop art style, colorful",
-        "Vintage": "sepia, 1950s film look, grain, nostalgia",
-        "Cyberpunk": "neon lights, cyberpunk, blue and pink hues, futuristic"
-    }
+    if fps == 0: fps = 25
     
-    full_prompt = f"{prompt}, {style_prompts.get(style_choice, '')}, colorized video"
-    negative_prompt = "black and white, monochrome, greyscale, low quality, blurry, distorted, nsfw, watermark"
-    
-    colored_frames = []
-    
-    print("🎬 Starting Frame Processing on ZeroGPU...")
+    frame_count = 0
+    print("🎬 Processing frames...")
     
     while True:
         ret, frame = cap.read()
         if not ret:
             break
             
-        # تحويل من BGR إلى RGB
-        pil_image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
-        canny_image = canny_processor(pil_image)
+        # DDColor يقبل الصورة بصيغة BGR أو RGB (مصفوفة Numpy)
+        # نقوم بتمرير الإطار مباشرة للنموذج
+        
+        # المعالجة باستخدام DDColor
+        # الناتج يكون قاموساً يحتوي على الصورة الملونة تحت مفتاح 'output_img'
+        result = ddcolor_pipeline(frame)
+        colorized_frame_bgr = result['output_img']
         
-        # التوليد
-        with torch.inference_mode():
-            image_out = pipe(
-                prompt=full_prompt,
-                negative_prompt=negative_prompt,
-                image=canny_image,
-                num_inference_steps=steps,
-                guidance_scale=1.5,
-                controlnet_conditioning_scale=1.0
-            ).images[0]
+        # حفظ الإطار كصورة PNG (لتجنب مشاكل ترميز الفيديو في OpenCV)
+        frame_filename = os.path.join(temp_frames_dir, f"frame_{frame_count:05d}.png")
+        cv2.imwrite(frame_filename, colorized_frame_bgr)
         
-        colored_frames.append(np.array(image_out))
+        frame_count += 1
+        if frame_count % 10 == 0:
+            print(f"Processed {frame_count} frames...")
     
     cap.release()
+    print(f"✅ Finished processing {frame_count} frames. Stitching video...")
 
-    # --- تجميع الفيديو ---
-    fourcc = cv2.VideoWriter_fourcc(*'mp4v') 
-    out = cv2.VideoWriter(output_temp_video_no_audio, fourcc, fps, (width, height))
+    # --- 3. تجميع الفيديو باستخدام FFmpeg ---
+    # استخدام نمط %05d لقراءة الإطارات بالترتيب الصحيح
+    input_frames = ffmpeg.input(os.path.join(temp_frames_dir, 'frame_%05d.png'), framerate=fps)
     
-    for frame in colored_frames:
-        out.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
-        
-    out.release()
-    
-    # --- دمج الصوت ---
     if audio_exists:
-        try:
-            (
-                ffmpeg
-                .input(output_temp_video_no_audio)
-                .output(ffmpeg.input(audio_path).audio, final_output_name, vcodec='copy', acodec='copy')
-                .run(overwrite_output=True, quiet=True)
-            )
-        except ffmpeg.Error:
-            shutil.copy(output_temp_video_no_audio, final_output_name)
+        input_audio = ffmpeg.input(audio_path)
+        # استخدام ترميز x264 لضمان التوافقية
+        stream = ffmpeg.output(input_frames, input_audio, final_output_name, vcodec='libx264', pix_fmt='yuv420p', acodec='aac', shortest=None)
     else:
-        shutil.copy(output_temp_video_no_audio, final_output_name)
-        
+        stream = ffmpeg.output(input_frames, final_output_name, vcodec='libx264', pix_fmt='yuv420p')
+    
+    try:
+        stream.run(overwrite_output=True, quiet=True)
+    except ffmpeg.Error as e:
+        print("FFmpeg Error:", e.stderr.decode('utf8'))
+        # محاولة أخيرة بدون صوت في حال فشل الدمج
+        ffmpeg.input(os.path.join(temp_frames_dir, 'frame_%05d.png'), framerate=fps).output(final_output_name, vcodec='libx264', pix_fmt='yuv420p').run(overwrite_output=True)
+
+    # تنظيف
+    shutil.rmtree(temp_frames_dir, ignore_errors=True)
     if os.path.exists(audio_path): os.remove(audio_path)
     
     return final_output_name
 
 # ==========================================
-# 3. واجهة Gradio
+# 3. واجهة التطبيق (بسيطة واحترافية)
 # ==========================================
 custom_css = """
 #col-container {max-width: 700px; margin-left: auto; margin-right: auto;}
 """
 
-with gr.Blocks(css=custom_css, title="Turbo Video Colorizer") as demo:
+with gr.Blocks(css=custom_css, title="Professional Video Colorizer (DDColor)") as demo:
     with gr.Column(elem_id="col-container"):
-        gr.Markdown("# ⚡ Turbo Video Colorizer (LCM + ZeroGPU)")
-        gr.Markdown("تلوين الفيديو بسرعة عالية باستخدام ZeroGPU و LCM-LoRA.")
+        gr.Markdown("# 🎞️ Professional Video Colorizer")
+        gr.Markdown("تلوين احترافي وواقعي للفيديو باستخدام نموذج DDColor. يحافظ على التفاصيل الأصلية بدون تغيير.")
         
         with gr.Row():
-            video_input = gr.Video(label="رفع الفيديو")
-        
-        with gr.Row():
-            prompt_input = gr.Textbox(label="وصف المشهد", placeholder="مثال: A sunny day in the park")
-            style_dropdown = gr.Dropdown(
-                ["Auto Color", "Vivid", "Vintage", "Cyberpunk"], 
-                label="النمط", value="Auto Color"
-            )
-        
-        steps_slider = gr.Slider(minimum=4, maximum=10, step=1, value=5, label="الخطوات (5 recommended)")
+            video_input = gr.Video(label="فيديو أبيض وأسود (Input)")
+            video_output = gr.Video(label="الفيديو الملون (Output)")
         
-        submit_btn = gr.Button("🎨 تلوين الفيديو", variant="primary")
-        video_output = gr.Video(label="النتيجة")
+        # لا توجد خيارات إضافية، فقط زر التلوين
+        submit_btn = gr.Button("✨ بدء التلوين الاحترافي", variant="primary", size="lg")
         
         submit_btn.click(
-            fn=colorize_video_multistyle,
-            inputs=[video_input, prompt_input, style_dropdown, steps_slider],
+            fn=colorize_video_professional,
+            inputs=[video_input],
             outputs=video_output
         )