app.py

# app.py
# Robust Hugging Face Space: load a Diffusers model with safe fallbacks
# No branding in source — ready to publish under any HF account

import os
import time
import traceback
import logging
from typing import Optional, Tuple

import torch
from PIL import Image
import gradio as gr

from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
from huggingface_hub import hf_hub_download, HfApi
from transformers import logging as trf_logging

# -------------------------
# Logging
# -------------------------
logging.basicConfig(level=logging.INFO, format="%(asctime)s — %(levelname)s — %(message)s")
logger = logging.getLogger("prompt-image-editor")
trf_logging.set_verbosity_error()

# -------------------------
# Config via environment
# -------------------------
MODEL_ID = os.getenv("MODEL_ID", "runwayml/stable-diffusion-v1-5")  # recommended default
HF_TOKEN = os.getenv("HF_API_TOKEN")  # optional, put as Secret if needed
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
RETRY_COUNT = int(os.getenv("MODEL_LOAD_RETRIES", "3"))
RETRY_WAIT_SECONDS = float(os.getenv("MODEL_LOAD_RETRY_WAIT", "2.0"))

# Optional: switch to inference-api mode instead of loading the model in-process
USE_INFERENCE_API = os.getenv("USE_INFERENCE_API", "false").lower() in ("1", "true", "yes")

# -------------------------
# Utilities
# -------------------------
def safe_from_pretrained(model_id: str, token: Optional[str] = None):
    """
    Load a diffusers pipeline with safe options (dtype/device_map when available).
    Raise exception to caller on failure.
    """
    kwargs = {}
    if token:
        kwargs["use_auth_token"] = token

    # Use float16 on CUDA for memory saving; else float32
    torch_dtype = torch.float16 if DEVICE == "cuda" else torch.float32

    # Try to create pipeline with recommended scheduler
    pipe = StableDiffusionPipeline.from_pretrained(
        model_id,
        torch_dtype=torch_dtype,
        **kwargs
    )

    # set scheduler if desired (optional improvement)
    try:
        pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
    except Exception:
        # ignore if incompatible
        pass

    if DEVICE == "cuda":
        pipe = pipe.to("cuda")
    else:
        pipe = pipe.to("cpu")
    # enable VAE tiling or RAM optimizations if needed (user can extend)
    return pipe

def load_pipeline_with_retries(model_id: str, token: Optional[str], retries: int = 3, wait: float = 2.0):
    """
    Attempt to load model with retries and fallback logic.
    Returns (pipeline or None, error_message or None)
    """
    last_err = None
    for attempt in range(1, retries + 1):
        try:
            logger.info(f"[load] Attempt {attempt}/{retries} to load model '{model_id}' (token set: {'yes' if token else 'no'}).")
            pipe = safe_from_pretrained(model_id, token)
            logger.info(f"[load] Successfully loaded model '{model_id}'.")
            return pipe, None
        except Exception as e:
            last_err = traceback.format_exc()
            logger.warning(f"[load] Failed attempt {attempt}: {e}")
            if attempt < retries:
                time.sleep(wait * attempt)  # exponential-ish backoff
    # fallback attempt to a known public model if initial failed
    fallback = "runwayml/stable-diffusion-v1-5"
    if model_id != fallback:
        try:
            logger.info(f"[load] Trying fallback model '{fallback}'.")
            pipe = safe_from_pretrained(fallback, None)
            logger.info(f"[load] Successfully loaded fallback '{fallback}'.")
            return pipe, None
        except Exception as e:
            last_err = traceback.format_exc()
            logger.error(f"[load] Fallback also failed: {e}")
    return None, last_err

# -------------------------
# Pipeline init
# -------------------------
pipe = None
load_error = None

if USE_INFERENCE_API:
    logger.info("Configured to use Inference API mode. The app will not load local models.")
else:
    try:
        pipe, load_error = load_pipeline_with_retries(MODEL_ID, HF_TOKEN, retries=RETRY_COUNT, wait=RETRY_WAIT_SECONDS)
    except Exception as e:
        pipe = None
        load_error = traceback.format_exc()
        logger.error("Unexpected error during model load:\n" + load_error)

# -------------------------
# Inference function
# -------------------------
def generate_image(prompt: str, steps: int = 28, guidance: float = 7.5) -> Tuple[Optional[Image.Image], str]:
    """
    Returns (PIL.Image or None, status message)
    """
    if USE_INFERENCE_API:
        return None, "Inference API mode enabled — implement API call flow or disable USE_INFERENCE_API."
    if pipe is None:
        return None, "Model is not loaded. Check Space Settings (MODEL_ID & HF_API_TOKEN) and server logs."
    if not prompt or not prompt.strip():
        return None, "Please enter a valid prompt."

    try:
        # autocast only on CUDA
        if DEVICE == "cuda":
            with torch.autocast("cuda"):
                out = pipe(prompt=prompt, guidance_scale=guidance, num_inference_steps=int(steps))
        else:
            out = pipe(prompt=prompt, guidance_scale=guidance, num_inference_steps=int(steps))
        img = out.images[0]
        return img, "OK"
    except Exception as e:
        logger.exception("Inference failed")
        return None, f"Inference error: {str(e)}"

# -------------------------
# Gradio UI
# -------------------------
title = "Prompt Image Editor"
description = "Generate or edit images using a Diffusers-compatible model. Configure MODEL_ID and HF_API_TOKEN in Settings → Variables & Secrets."

with gr.Blocks(title=title) as demo:
    gr.Markdown(f"# {title}")
    gr.Markdown(description)

    with gr.Row():
        with gr.Column(scale=2):
            prompt = gr.Textbox(lines=4, label="Prompt", placeholder="e.g. A portrait of an astronaut riding a horse, cinematic lighting")
            steps = gr.Slider(minimum=10, maximum=60, step=1, value=28, label="Steps")
            guidance = gr.Slider(minimum=1.0, maximum=20.0, step=0.1, value=7.5, label="Guidance scale")
            run_btn = gr.Button("Generate")
            status = gr.Textbox(label="Status", interactive=False, value="Model loaded." if pipe else "Model not loaded. Check settings.")
        with gr.Column(scale=3):
            out_img = gr.Image(label="Output image", type="pil")

    def _on_generate(prompt_text, steps_val, guidance_val):
        img, msg = generate_image(prompt_text, steps_val, guidance_val)
        return img, msg

    run_btn.click(_on_generate, inputs=[prompt, steps, guidance], outputs=[out_img, status])

if __name__ == "__main__":
    demo.launch()