import gradio as gr
import torch
from PIL import Image

print("🚀 جاري تحميل Pony Diffusion V6 XL (Uncensored)...")

try:
    from diffusers import StableDiffusionXLImg2ImgPipeline
    pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained(
        "LyliaEngine/Pony_Diffusion_V6_XL",
        torch_dtype=torch.float16,
        device_map="balanced",
        use_safetensors=True,
        variant="fp16"
    )
    pipe.enable_model_cpu_offload()   # ضروري للـ Free Tier
    print("✅ Pony Diffusion V6 XL محمل بنجاح - Uncensored")
except Exception as e:
    print(f"❌ {e}")
    pipe = None

def img2img_pony(input_image, prompt, strength=0.75, steps=30, seed=-1):
    if pipe is None:
        return None, "فشل تحميل النموذج"
    if input_image is None:
        return None, "ارفع صورة أولاً"
    
    try:
        generator = torch.Generator("cuda").manual_seed(seed) if seed != -1 else None
        
        # Pony يحتاج score خاص
        full_prompt = f"score_9, score_8_up, score_7_up, {prompt}"
        
        output = pipe(
            prompt=full_prompt,
            image=input_image,
            strength=strength,
            num_inference_steps=steps,
            generator=generator,
            guidance_scale=5.0,
        ).images[0]
        
        return output, "✅ تم التوليد بنجاح (غير مقيد)"
    except Exception as e:
        return None, f"خطأ: {str(e)}"


with gr.Blocks(title="Pony Diffusion Img2Img", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🖼️ Pony Diffusion V6 XL\n**Img2Img - Uncensored (غير مقيد)**")

    with gr.Row():
        with gr.Column():
            input_image = gr.Image(label="الصورة الأصلية", type="pil", height=512)
            prompt = gr.Textbox(
                label="تعليمات التعديل",
                lines=5,
                placeholder="change background to dark fantasy forest, add glowing mushrooms, cinematic lighting..."
            )
            
            with gr.Row():
                strength = gr.Slider(0.6, 0.95, value=0.75, label="Strength")
                steps = gr.Slider(20, 50, value=30, label="Steps")
            seed = gr.Number(-1, label="Seed (-1 = عشوائي)")

        with gr.Column():
            output_image = gr.Image(label="الصورة الناتجة", height=512)
            status = gr.Textbox(label="الحالة", interactive=False)

    btn = gr.Button("🚀 توليد (Uncensored)", variant="primary", size="large")
    btn.click(img2img_pony, [input_image, prompt, strength, steps, seed], [output_image, status])

demo.launch(server_name="0.0.0.0", server_port=7860, enable_queue=True)