import torch
from diffusers import StableDiffusionImg2ImgPipeline
import gradio as gr
from PIL import Image

# تحميل الموديل بصيغة مناسبة للـ CPU
pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=torch.float32
)
pipe = pipe.to("cpu")

def generate_image(input_image, prompt, strength=0.7, guidance=7.5):
    if input_image is None or prompt.strip() == "":
        return None
    
    # تغيير حجم الصورة لتسريع المعالجة
    init_image = input_image.convert("RGB").resize((512, 512))
    result = pipe(
        prompt=prompt,
        image=init_image,
        strength=strength,
        guidance_scale=guidance
    ).images[0]
    return result

# واجهة Gradio
with gr.Blocks() as demo:
    gr.Markdown("## 🖼️ Stable Diffusion Img2Img - ديكور من صورة")
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(type="pil", label="ارفع الصورة")
            prompt = gr.Textbox(label="الوصف (Prompt)", placeholder="مثال: غرفة مع أثاث عصري وإضاءة دافئة")
            strength = gr.Slider(0.1, 1.0, value=0.7, step=0.1, label="درجة التغيير (Strength)")
            guidance = gr.Slider(1, 15, value=7.5, step=0.5, label="جودة التفاصيل (Guidance Scale)")
            btn = gr.Button("توليد الصورة")
        with gr.Column():
            output_image = gr.Image(label="النتيجة")
    
    btn.click(
        fn=generate_image,
        inputs=[input_image, prompt, strength, guidance],
        outputs=output_image
    )

demo.launch()