Update app.py

app.py

@@ -1,38 +1,30 @@
 # ---- Flags ----
 run_api = False
-SSD_1B = 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
-from IPython.display import clear_output
 
-# ---- Minimal, deterministic env bootstrap (optional) ----
-# Prefer pinning in requirements.txt instead of installing here.
-def check_environment():
-    try:
-        import torch  # noqa: F401
-        print("Environment is already installed.")
-    except ImportError:
-        print("Environment not found. Installing pinned dependencies...")
-        # Strongly prefer doing this via requirements.txt at build time.
-        os.system("pip install --upgrade pip")
-        os.system("pip install diffusers==0.30.0 transformers>=4.41.0 accelerate>=0.31.0 huggingface_hub>=0.23.4 safetensors>=0.4.2 gradio==4.37.1 python-dotenv")
-        clear_output()
-        print("Environment installed successfully.")
-
-check_environment()
-
-# ---- App imports (safe after environment check) ----
+# 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
 
-# Optional: only imported if SSD_1B=True
-# from diffusers import AutoPipelineForText2Image
-
 # ---- Config / constants ----
 current_dir = os.getcwd()
 cache_path = os.path.join(current_dir, "cache")
@@ -42,31 +34,37 @@ 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 / NVML detection (robust) ----
+# ---- 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:
+        if proc.returncode == 0 and proc.stdout.strip():
             print(proc.stdout)
         else:
-            # Show the stderr to aid debugging, but don't trust it for logic
-            print(proc.stderr or "nvidia-smi returned a non-zero exit code.")
+            # 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}")
 
-# dtype & device
-dtype = torch.float16 if is_gpu else torch.float32
 device = torch.device("cuda") if is_gpu else torch.device("cpu")
+dtype = torch.float16 if is_gpu else torch.float32
 
-# Optional: fewer surprises when CUDA is flaky
-if not is_gpu:
-    # Avoid cuda-related env flags when no GPU
-    os.environ.pop("CUDA_LAUNCH_BLOCKING", None)
+# ---- 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:
@@ -74,26 +72,26 @@ if not SSD_1B:
     unet = UNet2DConditionModel.from_pretrained(
         "latent-consistency/lcm-sdxl",
         torch_dtype=dtype,
-        variant="fp16" if is_gpu else None,
         cache_dir=cache_path,
+        **_add_variant({})
     )
     pipe = DiffusionPipeline.from_pretrained(
         "stabilityai/stable-diffusion-xl-base-1.0",
         unet=unet,
         torch_dtype=dtype,
-        variant="fp16" if is_gpu else None,
         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  # local import
+    from diffusers import AutoPipelineForText2Image
     pipe = AutoPipelineForText2Image.from_pretrained(
         "segmind/SSD-1B",
         torch_dtype=dtype,
-        variant="fp16" if is_gpu else None,
         cache_dir=cache_path,
+        **_add_variant({})
     )
     pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
     pipe.to(device)
@@ -111,12 +109,15 @@ def generate(
     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.")
-    # Make sure sizes are sane on CPU
+
+    # 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))
@@ -133,8 +134,6 @@ def generate(
     )
     return out.images[0]
 
-clear_output()
-
 # ---- Optional notebook helper ----
 def generate_image(prompt="A scenic watercolor landscape, mountains at dawn"):
     img = generate(
@@ -147,10 +146,13 @@ def generate_image(prompt="A scenic watercolor landscape, mountains at dawn"):
         num_inference_steps=4,
         secret_token=SECRET_TOKEN,
     )
-    from IPython.display import display
-    display(img)
+    try:
+        from IPython.display import display
+        display(img)
+    except Exception:
+        pass  # Non-notebook environment
 
-# ---- UI ----
+# ---- UI (Gradio 3.39.0 components) ----
 if not run_api:
     secret_token = gr.Textbox(
         label="Secret Token",