app.py

# ---- Flags ----
run_api = False
SSD_1B = False  # True = use SSD-1B + LCM LoRA, False = SDXL Base + LCM (default)

# ---- Standard imports ----
import os
import subprocess
import numpy as np

# Optional: clear_output is nice in notebooks; ignore if not available
try:
    from IPython.display import clear_output  # noqa: F401
except Exception:
    def clear_output():  # no-op outside notebooks
        pass

# ---- Tame NVML noise in containers without GPU drivers (optional) ----
os.environ.setdefault("DEEPSPEED_DISABLE_NVML", "1")
import warnings
warnings.filterwarnings("ignore", message="Can't initialize NVML")

# ---- App imports (expect deps from requirements.txt already installed) ----
import torch
import gradio as gr
from PIL import Image
from diffusers import UNet2DConditionModel, DiffusionPipeline, LCMScheduler

# ---- Config / constants ----
current_dir = os.getcwd()
cache_path = os.path.join(current_dir, "cache")
os.makedirs(cache_path, exist_ok=True)

MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "1024"))
SECRET_TOKEN = os.getenv("SECRET_TOKEN", "default_secret")

# ---- GPU visibility / info (for logs only) ----
def print_nvidia_smi():
    try:
        proc = subprocess.run(["nvidia-smi"], capture_output=True, text=True)
        if proc.returncode == 0 and proc.stdout.strip():
            print(proc.stdout)
        else:
            # Show stderr when present to help debugging; not used for logic
            if proc.stderr:
                print(proc.stderr)
            else:
                print("nvidia-smi not available or returned no output.")
    except FileNotFoundError:
        print("nvidia-smi not found on PATH.")

print_nvidia_smi()

# ---- Device + dtype selection (robust) ----
is_gpu = torch.cuda.is_available()
print(f"CUDA available: {is_gpu}")

device = torch.device("cuda") if is_gpu else torch.device("cpu")
dtype = torch.float16 if is_gpu else torch.float32

# ---- Helpers to only pass 'variant' when needed (Diffusers <=0.23 friendly) ----
def _add_variant(kwargs: dict) -> dict:
    """Only include 'variant' when running on GPU."""
    if is_gpu:
        kwargs = dict(kwargs)  # shallow copy
        kwargs["variant"] = "fp16"
    return kwargs

# ---- Pipeline setup ----
if not SSD_1B:
    # SDXL base + LCM UNet
    unet = UNet2DConditionModel.from_pretrained(
        "latent-consistency/lcm-sdxl",
        torch_dtype=dtype,
        cache_dir=cache_path,
        **_add_variant({})
    )
    pipe = DiffusionPipeline.from_pretrained(
        "stabilityai/stable-diffusion-xl-base-1.0",
        unet=unet,
        torch_dtype=dtype,
        cache_dir=cache_path,
        **_add_variant({})
    )
    pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
    pipe.to(device)
else:
    # SSD-1B + LCM LoRA
    from diffusers import AutoPipelineForText2Image
    pipe = AutoPipelineForText2Image.from_pretrained(
        "segmind/SSD-1B",
        torch_dtype=dtype,
        cache_dir=cache_path,
        **_add_variant({})
    )
    pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
    pipe.to(device)
    pipe.load_lora_weights("latent-consistency/lcm-lora-ssd-1b")
    pipe.fuse_lora()

# ---- Core generate function ----
def generate(
    prompt: str,
    negative_prompt: str = "",
    seed: int = 0,
    width: int = 1024,
    height: int = 1024,
    guidance_scale: float = 0.0,
    num_inference_steps: int = 4,
    secret_token: str = "",
) -> Image.Image:
    # Token gate
    if secret_token != SECRET_TOKEN:
        raise gr.Error("Invalid secret token. Set SECRET_TOKEN on the server or pass the correct token.")

    # Clamp sizes (avoid OOM on CPU)
    width = int(np.clip(width, 256, MAX_IMAGE_SIZE))
    height = int(np.clip(height, 256, MAX_IMAGE_SIZE))

    # Deterministic generator on the active device
    generator = torch.Generator(device=device)
    if seed is not None:
        generator = generator.manual_seed(int(seed))

    out = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        width=width,
        height=height,
        guidance_scale=guidance_scale,
        num_inference_steps=num_inference_steps,
        generator=generator,
        output_type="pil",
    )
    return out.images[0]

# ---- Optional notebook helper ----
def generate_image(prompt="A scenic watercolor landscape, mountains at dawn"):
    img = generate(
        prompt=prompt,
        negative_prompt="",
        seed=0,
        width=1024,
        height=1024,
        guidance_scale=0.0,
        num_inference_steps=4,
        secret_token=SECRET_TOKEN,
    )
    try:
        from IPython.display import display
        display(img)
    except Exception:
        pass  # Non-notebook environment

# ---- UI (Gradio 3.39.0 components) ----
if not run_api:
    secret_token = gr.Textbox(
        label="Secret Token",
        placeholder="Enter your secret token",
        type="password",
    )
    prompt = gr.Textbox(
        label="Prompt",
        show_label=True,
        max_lines=2,
        placeholder="Enter your prompt",
    )
    negative_prompt = gr.Textbox(
        label="Negative prompt",
        max_lines=2,
        placeholder="Enter a negative prompt (optional)",
    )
    seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
    width = gr.Slider(label="Width", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=1024)
    height = gr.Slider(label="Height", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=1024)
    guidance_scale = gr.Slider(label="Guidance scale", minimum=0, maximum=2, step=0.1, value=0.0)
    num_inference_steps = gr.Slider(label="Inference steps", minimum=1, maximum=8, step=1, value=4)

    iface = gr.Interface(
        fn=generate,
        inputs=[prompt, negative_prompt, seed, width, height, guidance_scale, num_inference_steps, secret_token],
        outputs=gr.Image(label="Result"),
        title="Image Generator (LCM)",
        description="Fast SDXL/SSD-1B image generation with LCM. Uses CPU if CUDA is unavailable.",
    )
    iface.launch()

if run_api:
    with gr.Blocks() as demo:
        gr.Markdown(
            "### REST API for LCM Text-to-Image\n"
            "Use the `/run` endpoint programmatically with your secret."
        )
        secret_token = gr.Textbox(label="Secret Token", type="password")
        prompt = gr.Textbox(label="Prompt")
        negative_prompt = gr.Textbox(label="Negative prompt")
        seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
        width = gr.Slider(label="Width", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=1024)
        height = gr.Slider(label="Height", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=1024)
        guidance_scale = gr.Slider(label="Guidance scale", minimum=0, maximum=2, step=0.1, value=0.0)
        num_inference_steps = gr.Slider(label="Inference steps", minimum=1, maximum=8, step=1, value=4)

        inputs = [prompt, negative_prompt, seed, width, height, guidance_scale, num_inference_steps, secret_token]
        prompt.submit(fn=generate, inputs=inputs, outputs=gr.Image(), api_name="run")

    demo.queue(max_size=32).launch(debug=False)