import gradio as gr
from diffusers import StableDiffusionImg2ImgPipeline, DDIMScheduler
import torch
from PIL import Image
import time
import os

print("=" * 50)
print("🔄 در حال لود مدل...")
print("=" * 50)

# تشخیص دستگاه
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float16 if device == "cuda" else torch.float32

print(f"📍 دستگاه: {device}")

# لود مدل
pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
    "hafsa000/interior-design",
    torch_dtype=dtype,
    safety_checker=None
)
pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
pipe = pipe.to(device)

# بهینه‌سازی برای CPU
if device == "cpu":
    pipe.enable_attention_slicing()
    print("⚠️ حالت CPU: پردازش کند خواهد بود (5-15 دقیقه هر تصویر)")
else:
    print("✅ مدل روی GPU لود شد!")

# استایل‌ها
STYLES = {
    "Minimalistic": "minimalist interior design, clean lines, neutral colors, natural light, photorealistic, 8k",
    "Boho": "bohemian interior design, colorful textiles, plants, eclectic decor, photorealistic, 8k",
    "Farmhouse": "farmhouse interior design, rustic wood, vintage decor, warm colors, photorealistic",
    "Neoclassical": "neoclassical interior design, elegant, marble, gold accents, photorealistic, 8k",
    "Parisian": "parisian apartment interior, chic decor, elegant furniture, photorealistic, 8k",
    "Scandinavian": "scandinavian interior design, white walls, light wood, cozy textiles, plants, photorealistic",
    "Japanese": "japanese interior design, minimal, zen atmosphere, natural materials, photorealistic",
    "Midcentury Modern": "midcentury modern interior design, retro furniture, organic shapes, warm wood, photorealistic",
    "Industrial": "industrial style interior, exposed brick, metal accents, raw materials, photorealistic",
    "Luxury": "luxury interior design, marble, gold accents, premium materials, photorealistic, 8k",
}

def redesign(image, style_name, strength, guidance):
    if image is None:
        return None, "❌ لطفاً عکس آپلود کنید"
    
    print(f"🎨 پردازش با سبک: {style_name}")
    start_time = time.time()
    
    prompt = STYLES.get(style_name, STYLES["Minimalistic"])
    
    try:
        # تنظیمات برای CPU (اگه CPU هست، steps کمتر)
        num_steps = 15 if device == "cpu" else 30
        
        result = pipe(
            prompt=prompt,
            image=image,
            strength=float(strength),
            guidance_scale=float(guidance),
            negative_prompt="blurry, low quality, distorted, cartoon, drawing",
            num_inference_steps=num_steps
        ).images[0]
        
        elapsed = time.time() - start_time
        print(f"✅ تمام شد در {elapsed:.1f} ثانیه")
        return result, f"✅ طراحی کامل شد ({elapsed:.1f} ثانیه)"
    
    except Exception as e:
        print(f"❌ خطا: {e}")
        import traceback
        traceback.print_exc()
        return None, f"❌ خطا: {str(e)}"

# رابط کاربری
with gr.Blocks(title="AI Interior Designer", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🏠 طراحی اتاق با هوش مصنوعی")
    gr.Markdown(f"⚙️ در حال اجرا روی: **{device.upper()}**")
    
    with gr.Row():
        with gr.Column():
            img_input = gr.Image(type="pil", label="📷 عکس اتاق", height=400)
            style = gr.Dropdown(
                choices=list(STYLES.keys()),
                value="Minimalistic",
                label="🎨 سبک طراحی"
            )
            with gr.Row():
                strength = gr.Slider(0.3, 0.9, value=0.65, step=0.05, label="💪 میزان تغییر")
                guidance = gr.Slider(5.0, 15.0, value=7.5, step=0.5, label="🎯 دقت به پرامپت")
            btn = gr.Button("✨ طراحی کن!", variant="primary", size="lg")
            status = gr.Textbox(label="وضعیت", interactive=False)
        
        with gr.Column():
            output = gr.Image(label="🖼️ نتیجه طراحی", height=400)
    
    btn.click(redesign, [img_input, style, strength, guidance], [output, status])
    

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