import torch
from diffusers import StableDiffusionXLPipeline, DPMSolverMultistepScheduler
import gradio as gr
from openai import OpenAI
import os
import spaces # المكتبة السحرية للـ GPU المجاني

# 1. إعداد العميل (تأكد إنك ضفت المفتاح في الـ Secrets باسم OPENROUTER_API_KEY)
api_key = os.getenv("OPENROUTER_API_KEY")
client = OpenAI(base_url="https://openrouter.ai/api/v1", api_key=api_key)

# 2. تحميل الموديل (بيتحمل على الـ CPU في الأول وده عادي)
model_id = "SG161222/RealVisXL_V5.0_Lightning"
pipe = StableDiffusionXLPipeline.from_pretrained(
    model_id, 
    torch_dtype=torch.float16, 
    variant="fp16",
    use_safetensors=True
)
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config, use_karras_sigmas=True)

# 3. دالة التوليد (لاحظ سطر @spaces.GPU)
@spaces.GPU(duration=60) 
def generate_image(user_idea):
    # السطر ده هو السر: بينقل الموديل للكارت "وقت الضغط على الزرار" فقط
    pipe.to("cuda") 
    
    system_rules = "Expert in cinematic photography. Transform idea to 8k advertising prompt. Return ONLY English."
    
    try:
        response = client.chat.completions.create(
            model="google/gemini-2.0-flash-001",
            messages=[{"role": "system", "content": system_rules}, {"role": "user", "content": user_idea}]
        )
        hyper_prompt = response.choices[0].message.content.strip()
    except:
        hyper_prompt = user_idea

    # توليد الصورة
    image = pipe(
        prompt=hyper_prompt,
        negative_prompt="cartoon, anime, plastic, blurry, lowres, text, watermark",
        num_inference_steps=7, 
        guidance_scale=2.0
    ).images[0]
    
    return image, hyper_prompt

# 4. واجهة Gradio
with gr.Blocks() as demo:
    gr.Markdown("# 💎 Luxury AI Studio")
    with gr.Row():
        with gr.Column():
            inp = gr.Textbox(label="أوصف فكرتك بالبلدي")
            btn = gr.Button("توليد التصميم ✨", variant="primary")
        with gr.Column():
            out_img = gr.Image(label="النتيجة")
            out_prompt = gr.Markdown()

    btn.click(generate_image, inputs=[inp], outputs=[out_img, out_prompt])

demo.launch()