import os
import gc
import time
import random
import torch
import gradio as gr

from diffusers import DiffusionPipeline

# =========================
# HARD CPU MODE
# =========================
os.environ["CUDA_VISIBLE_DEVICES"] = ""
os.environ["HF_HUB_DISABLE_TELEMETRY"] = "1"
os.environ["TOKENIZERS_PARALLELISM"] = "false"

cpu_cores = os.cpu_count() or 1
torch.set_num_threads(cpu_cores)
torch.set_num_interop_threads(cpu_cores)

os.environ["OMP_NUM_THREADS"] = str(cpu_cores)
os.environ["MKL_NUM_THREADS"] = str(cpu_cores)

torch.backends.mkldnn.enabled = True

device = torch.device("cpu")
dtype = torch.bfloat16 if torch.cpu.is_bf16_supported() else torch.float32

MODEL_ID = "tensorart/stable-diffusion-3.5-medium-turbo"
CACHE_DIR = "models"

# =========================
# LOAD PIPELINE
# =========================
def load_pipeline():
    pipe = DiffusionPipeline.from_pretrained(
        MODEL_ID,
        torch_dtype=dtype,
        cache_dir=CACHE_DIR,
        low_cpu_mem_usage=True
    )

    pipe.enable_attention_slicing()
    pipe.enable_vae_slicing()
    pipe.enable_sequential_cpu_offload()

    pipe = pipe.to(device)

    return pipe

pipe = load_pipeline()

# =========================
# GENERATION
# =========================
def generate(prompt, seed, progress=gr.Progress()):
    if not prompt:
        raise gr.Error("Prompt required")

    if seed < 0:
        seed = random.randint(0, 2**31 - 1)

    generator = torch.Generator(device=device).manual_seed(seed)

    steps = 6
    width = 512
    height = 512

    start = time.time()

    def callback(step, timestep, latents):
        done = step + 1
        elapsed = time.time() - start
        eta = (elapsed / done) * (steps - done)
        progress(done / steps, desc=f"Step {done}/{steps} | ETA {eta:.1f}s")

    with torch.inference_mode():
        gc.collect()
        image = pipe(
            prompt=prompt,
            width=width,
            height=height,
            num_inference_steps=steps,
            guidance_scale=0.0,
            generator=generator,
            callback=callback,
            callback_steps=1
        ).images[0]
        gc.collect()

    return image, seed

# =========================
# UI
# =========================
with gr.Blocks(title="SD 3.5 Medium Turbo CPU Ultra Lean") as demo:
    gr.Markdown("# Stable Diffusion 3.5 Medium Turbo — 16GB CPU Mode")

    prompt = gr.Textbox(label="Prompt", lines=3)
    seed = gr.Number(label="Seed (-1 random)", value=-1, precision=0)
    btn = gr.Button("Generate")

    image_out = gr.Image()
    seed_out = gr.Number(interactive=False)

    btn.click(generate, inputs=[prompt, seed], outputs=[image_out, seed_out])

demo.queue(max_size=5, concurrency_count=1)

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