# app.py — SDXL + LoRA from Hub (PEFT) with version-safe adapter weighting
import torch, gradio as gr
from diffusers import StableDiffusionXLPipeline, DPMSolverMultistepScheduler

# --- Config ---
BASE_MODEL = "stabilityai/stable-diffusion-xl-base-1.0"
LORA_REPO  = "freshcodestech/LingoSpace"      # your HF LoRA repo
LORA_FILE  = "LingoSpace-10.safetensors"      # exact filename in that repo
ADAPTER    = "lingospace"                     # local nickname; any safe string
LORA_W     = 0.8                              # 0.6–1.0 typical

# Device / dtype
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype  = torch.float16 if device == "cuda" else torch.float32

# --- Load SDXL base (disable watermarker to avoid extra deps on Spaces) ---
pipe = StableDiffusionXLPipeline.from_pretrained(
    BASE_MODEL,
    torch_dtype=dtype,
    use_safetensors=True,
    add_watermarker=False,
    **({"variant": "fp16"} if device == "cuda" else {})  # fp16 only on GPU
)
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
pipe = pipe.to(device)

# --- Load LoRA from Hub via PEFT backend ---
pipe.load_lora_weights(
    LORA_REPO,
    weight_name=LORA_FILE,
    adapter_name=ADAPTER,
    use_peft_backend=True
)

# Try all known set_adapters signatures; else fall back to per-call scaling
USE_RUNTIME_SCALE = False
if hasattr(pipe, "set_adapters"):
    try:
        # Newer diffusers: keyword
        pipe.set_adapters(ADAPTER, adapter_weights=LORA_W)
    except TypeError:
        try:
            # Older diffusers: 2 positional lists
            pipe.set_adapters([ADAPTER], [LORA_W])
        except Exception:
            USE_RUNTIME_SCALE = True
else:
    USE_RUNTIME_SCALE = True

# Memory helpers (T4)
if device == "cuda":
    if hasattr(pipe, "enable_xformers_memory_efficient_attention"):
        pipe.enable_xformers_memory_efficient_attention()
    pipe.enable_attention_slicing()
    pipe.enable_vae_tiling()

# Defaults (smaller on CPU)
DEF_STEPS = 30 if device == "cuda" else 15
DEF_H = 1024 if device == "cuda" else 768
DEF_W = 1024 if device == "cuda" else 768

def infer(prompt, steps, cfg, h, w, lora_strength):
    # If we couldn’t set adapter weight globally, pass strength per call
    kwargs = {}
    if USE_RUNTIME_SCALE:
        kwargs["cross_attention_kwargs"] = {"scale": float(lora_strength)}
    else:
        # If API exists, update weight dynamically per request
        try:
            pipe.set_adapters(ADAPTER, adapter_weights=float(lora_strength))
        except TypeError:
            pipe.set_adapters([ADAPTER], [float(lora_strength)])

    out = pipe(
        prompt,
        num_inference_steps=int(steps),
        guidance_scale=float(cfg),
        height=int(h),
        width=int(w),
        **kwargs
    )
    return out.images[0]

DEFAULT_PROMPT = (
    "lingospace, cartoon, 2d illustration, flat colors, clean lineart, "
    "a raccoon in a compact utility space suit, holding a steaming bowl of rice, "
    "stars and soft nebula, cozy composition"
)

demo = gr.Interface(
    fn=infer,
    inputs=[
        gr.Textbox(label="Prompt", value=DEFAULT_PROMPT, lines=3),
        gr.Slider(10, 50, value=DEF_STEPS, step=1, label="Steps"),
        gr.Slider(3.0, 9.0, value=7.5, step=0.5, label="CFG"),
        gr.Slider(512, 1216, value=DEF_H, step=64, label="Height"),
        gr.Slider(512, 1216, value=DEF_W, step=64, label="Width"),
        gr.Slider(0.3, 1.2, value=LORA_W, step=0.05, label="LoRA Strength"),
    ],
    outputs=gr.Image(label="Output"),
    title="SDXL + LoRA (freshcodestech/LingoSpace)",
    description=f"Device: {device.upper()} • Base: {BASE_MODEL} • LoRA: {LORA_REPO}/{LORA_FILE}"
)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860, show_api=True)