Update app.py

app.py

@@ -1,110 +1,174 @@
 # 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 gradio as gr
-from PIL import Image
+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
 
-from diffusers import StableDiffusionPipeline
-from transformers import logging
-logging.set_verbosity_error()
+# -------------------------
+# 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 from environment
-MODEL_ID = os.getenv("MODEL_ID", "stabilityai/stable-diffusion-2-1")
-HF_TOKEN = os.getenv("HF_API_TOKEN")  # Secret in Spaces (optional)
+# -------------------------
+# 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"))
 
-def try_load(model_id, token=None):
+# 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):
     """
-    Try to load a diffusers pipeline. Raises the original exception on fatal error.
+    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
-    # choose dtype based on device
+
+    # 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 = StableDiffusionPipeline.from_pretrained(
-            model_id,
-            revision="fp16" if DEVICE == "cuda" else None,
-            torch_dtype=torch_dtype,
-            **kwargs
-        )
-        if DEVICE == "cuda":
-            pipe = pipe.to("cuda")
-        else:
-            pipe = pipe.to("cpu")
-        return pipe
+        pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
     except Exception:
-        # re-raise to let caller decide (we'll handle fallback outside)
-        raise
+        # ignore if incompatible
+        pass
 
-# Load pipeline with fallback logic
-def load_pipeline_with_fallback():
-    tried = []
-    # first attempt: user-provided MODEL_ID with token if exists
-    try:
-        print(f"Attempting to load MODEL_ID='{MODEL_ID}' (token set: {'yes' if HF_TOKEN else 'no'}) on {DEVICE}")
-        return try_load(MODEL_ID, token=HF_TOKEN)
-    except Exception as e:
-        tried.append((MODEL_ID, str(e)))
-        print(f"Failed to load {MODEL_ID}: {e}")
+    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
 
-    # fallback: try a known-public model
-    fallback_model = "runwayml/stable-diffusion-v1-5"
+# -------------------------
+# 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:
-        print(f"Attempting fallback model '{fallback_model}'")
-        return try_load(fallback_model, token=None)
+        pipe, load_error = load_pipeline_with_retries(MODEL_ID, HF_TOKEN, retries=RETRY_COUNT, wait=RETRY_WAIT_SECONDS)
     except Exception as e:
-        tried.append((fallback_model, str(e)))
-        print(f"Failed to load fallback {fallback_model}: {e}")
-
-    # if we get here, nothing could be loaded
-    msg = "Failed to load any model. Tried: " + ", ".join([f"{m}: {err[:80]}" for m,err in tried])
-    raise RuntimeError(msg)
-
-# initialize
-try:
-    pipe = load_pipeline_with_fallback()
-except Exception as e:
-    # If pipeline can't be loaded, set pipe = None and keep running (UI will show error)
-    pipe = None
-    load_error = traceback.format_exc()
-    print("MODEL LOAD ERROR:\n", load_error)
+        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):
+# -------------------------
+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 not loaded. Check Space Settings (MODEL_ID & HF_API_TOKEN). See server logs."
+        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 provide a prompt."
+        return None, "Please enter a valid prompt."
+
     try:
-        with torch.autocast("cuda") if DEVICE == "cuda" else torch.no_grad():
-            out = pipe(prompt=prompt, guidance_scale=guidance, num_inference_steps=steps)
-            img = out.images[0]
+        # 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:
-        print("Inference error:", e)
+        logger.exception("Inference failed")
         return None, f"Inference error: {str(e)}"
 
+# -------------------------
 # Gradio UI
-with gr.Blocks(title="Prompt Image Editor") as demo:
-    gr.Markdown("# Prompt Image Editor")
+# -------------------------
+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():
-            prompt = gr.Textbox(lines=3, label="Prompt")
+        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.5, value=7.5, label="Guidance")
-            run = gr.Button("Generate")
-            status = gr.Textbox(label="Status")
-        with gr.Column():
-            out_img = gr.Image(label="Output", type="pil")
-
-    def _run(prompt, steps, guidance):
-        img, msg = generate_image(prompt, steps, guidance)
+            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.click(_run, inputs=[prompt, steps, guidance], outputs=[out_img, status])
+    run_btn.click(_on_generate, inputs=[prompt, steps, guidance], outputs=[out_img, status])
 
 if __name__ == "__main__":
     demo.launch()