import gradio as gr
import torch
from diffusers import AutoencoderKL
import numpy as np
from PIL import Image
import io
import base64

# تحميل نموذج VAE من Stable Diffusion
vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse", torch_dtype=torch.float16)
vae = vae.to("cuda" if torch.cuda.is_available() else "cpu")

def image_to_latent(image):
    # تحويل الصورة إلى التنسيق المناسب
    image = image.resize((512, 512))
    image = np.array(image).astype(np.float32) / 255.0
    image = image[None].transpose(0, 3, 1, 2)
    image = torch.from_numpy(image).to(vae.device)
    
    # استخراج التمثيل الكامن
    with torch.no_grad():
        latent = vae.encode(image).latent_dist.sample() * 0.18215
    
    return latent

def latent_to_image(latent):
    # إعادة بناء الصورة من التمثيل الكامن
    with torch.no_grad():
        image = vae.decode(latent / 0.18215).sample
    
    image = (image / 2 + 0.5).clamp(0, 1)
    image = image.cpu().permute(0, 2, 3, 1).numpy()
    image = (image[0] * 255).astype(np.uint8)
    image = Image.fromarray(image)
    
    return image

def apply_latent_poisoning(latent, strength=0.3):
    # إنشاء نسخة من التمثيل الكامن
    poisoned_latent = latent.clone()
    
    # توليد نمط تشويش مدروس
    noise = torch.randn_like(poisoned_latent) * strength
    
    # الحفاظ على البنية الرئيسية (مرشح تنعيم)
    noise = torch.nn.functional.avg_pool2d(noise, kernel_size=3, stride=1, padding=1)
    
    # إضافة الضوضاء إلى التمثيل الكامن
    poisoned_latent = poisoned_latent + noise
    
    return poisoned_latent

def process_image(input_image, strength=0.3):
    if input_image is None:
        return None
        
    # استخراج التمثيل الكامن
    latent = image_to_latent(input_image)
    
    # تطبيق التشويش على الفضاء الكامن
    poisoned_latent = apply_latent_poisoning(latent, float(strength))
    
    # إعادة بناء الصورة
    protected_image = latent_to_image(poisoned_latent)
    
    return protected_image

# تعريف واجهة Gradio الرئيسية
iface = gr.Interface(
    fn=process_image,
    inputs=[
        gr.Image(type="pil", label="الصورة الأصلية"),
        gr.Slider(minimum=0.1, maximum=0.5, value=0.3, step=0.1, label="قوة الحماية")
    ],
    outputs=gr.Image(type="pil", label="الصورة المحمية"),
    title="حماية الصور باستخدام الفضاء الكامن المعكوس",
    description="قم برفع صورة وتحديد قوة الحماية لتطبيق تقنية الفضاء الكامن المعكوس"
)

# تشغيل التطبيق
iface.launch( share=True)