Update app.py

app.py

@@ -1,19 +1,22 @@
 import os
 import uuid
 import random
-from typing import Tuple
+from typing import Tuple, Optional
 
 import gradio as gr
 import numpy as np
 from PIL import Image
 import torch
 import spaces
-from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
+from diffusers import (
+    StableDiffusionXLPipeline,
+    StableDiffusionPipeline,
+    EulerAncestralDiscreteScheduler,
+)
 
+PRIMARY_MODEL_ID = "SG161222/RealVisXL_V5.0_Lightning"   # requires access + token
+FALLBACK_MODEL_ID = "stabilityai/sd-turbo"               # public, fast 1.5-turbo
 
-# -----------------------
-# Style handling
-# -----------------------
 def apply_style(style_name: str, positive: str, negative: str = "") -> Tuple[str, str]:
     styles = {
         "3840 x 2160": (
@@ -25,67 +28,85 @@ def apply_style(style_name: str, positive: str, negative: str = "") -> Tuple[str
     }
     DEFAULT_STYLE_NAME = "3840 x 2160"
     p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
-    return p.replace("{prompt}", positive), (n + (" " + negative if negative else "")).strip()
+    n2 = (n + (" " + negative if negative else "")).strip()
+    return p.replace("{prompt}", positive), n2
 
+def _enable_performance_knobs():
+    if torch.cuda.is_available():
+        torch.backends.cuda.matmul.allow_tf32 = True
+        torch.set_grad_enabled(False)
 
-# -----------------------
-# Model loader
-# -----------------------
-def load_and_prepare_model() -> StableDiffusionXLPipeline:
-    model_id = "SG161222/RealVisXL_V5.0_Lightning"
+def _try_load_realvis(hf_token: Optional[str]):
     use_cuda = torch.cuda.is_available()
     dtype = torch.float16 if use_cuda else torch.float32
-    device = torch.device("cuda:0" if use_cuda else "cpu")
-
     pipe = StableDiffusionXLPipeline.from_pretrained(
-        model_id,
+        PRIMARY_MODEL_ID,
         torch_dtype=dtype,
         use_safetensors=True,
         add_watermarker=False,
+        token=hf_token,             # <- IMPORTANT
     )
-
-    # Use a stable, fast scheduler
     pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
-
-    # Memory/throughput optimizations
     try:
         pipe.enable_xformers_memory_efficient_attention()
     except Exception:
-        # xformers may be missing on CPU or some build types — that's fine
         pass
+    device = torch.device("cuda:0" if use_cuda else "cpu")
+    pipe = pipe.to(device)
+    return pipe
 
-    # Optional, safe perf knobs on CUDA
-    if use_cuda:
-        torch.backends.cuda.matmul.allow_tf32 = True
-        torch.set_grad_enabled(False)
-
+def _try_load_fallback():
+    # sd-turbo is Stable Diffusion 1.5 turbo; quick & public
+    use_cuda = torch.cuda.is_available()
+    dtype = torch.float16 if use_cuda else torch.float32
+    pipe = StableDiffusionPipeline.from_pretrained(
+        FALLBACK_MODEL_ID,
+        torch_dtype=dtype,
+        use_safetensors=True,
+    )
+    try:
+        pipe.enable_xformers_memory_efficient_attention()
+    except Exception:
+        pass
+    device = torch.device("cuda:0" if use_cuda else "cpu")
     pipe = pipe.to(device)
     return pipe
 
+def load_and_prepare_model():
+    _enable_performance_knobs()
+    hf_token = os.getenv("HF_TOKEN", "").strip() or None
 
-# Global model (loaded once per Space instance)
-model = load_and_prepare_model()
+    # Try RealVis first
+    try:
+        return _try_load_realvis(hf_token)
+    except Exception as e:
+        msg = str(e).lower()
+        if ("401" in msg or "403" in msg or "unauthorized" in msg or "forbidden" in msg):
+            # Clear hint in server logs; UI will still work via fallback.
+            print(
+                "\n[WARNING] Could not load RealVisXL (auth). "
+                "Make sure you've requested access and set HF_TOKEN in Space secrets.\n"
+            )
+        else:
+            print(f"\n[WARNING] RealVisXL failed to load: {e}\n")
+
+    # Fallback to sd-turbo so app still runs
+    print("[INFO] Falling back to stabilityai/sd-turbo (public).")
+    return _try_load_fallback()
 
+# Load once
+model = load_and_prepare_model()
 
-# -----------------------
-# Utils
-# -----------------------
 def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
     if randomize_seed:
         seed = random.randint(0, np.iinfo(np.int32).max)
     return int(seed)
 
-
 def save_image(img: Image.Image) -> str:
-    # Save to the working dir so HF can expose it as an artifact
     unique_name = f"{uuid.uuid4().hex}.png"
     img.save(unique_name)
     return unique_name
 
-
-# -----------------------
-# Generation
-# -----------------------
 @spaces.GPU(duration=60, enable_queue=True)
 def generate(
     prompt: str,
@@ -96,40 +117,36 @@ def generate(
     num_inference_steps: int = 25,
     randomize_seed: bool = False,
 ):
-    global model
-
-    # Guardrails
     if not prompt or not prompt.strip():
         raise gr.Error("Please enter a prompt.")
 
-    # SDXL prefers dims divisible by 8
+    # Make dimensions friendly for SD models
     width = int(max(256, (width // 8) * 8))
     height = int(max(256, (height // 8) * 8))
 
     seed = randomize_seed_fn(seed, randomize_seed)
     generator = torch.Generator(device=model.device).manual_seed(seed)
 
+    # If model is SDXL pipeline, use the styled prompts; if fallback SD1.5 turbo, style still OK
     positive_prompt, negative_prompt = apply_style("3840 x 2160", prompt)
 
-    # NOTE: pass strings (not one-element lists)
-    images = model(
+    # Some pipelines (sd-turbo) ignore guidance/steps or behave differently; passing is still safe
+    out = model(
         prompt=positive_prompt,
         negative_prompt=negative_prompt,
-        width=width,
-        height=height,
+        width=width if "xl" in model.__class__.__name__.lower() else None,
+        height=height if "xl" in model.__class__.__name__.lower() else None,
         guidance_scale=float(guidance_scale),
         num_inference_steps=int(num_inference_steps),
         generator=generator,
         output_type="pil",
-    ).images
+    )
 
+    # Handle both diffusers return shapes
+    images = getattr(out, "images", out)
     image_path = save_image(images[0])
     return image_path
 
-
-# -----------------------
-# UI
-# -----------------------
 with gr.Blocks(theme="soft") as demo:
     with gr.Row():
         with gr.Column(scale=12, elem_id="title_block"):
@@ -137,7 +154,7 @@ with gr.Blocks(theme="soft") as demo:
                 "<h1 style='text-align:center; color:white; font-weight:bold; text-decoration:underline;'>SNAPSCRIBE</h1>"
             )
             gr.Markdown(
-                "<h2 style='text-align:center; color:white; font-weight:bold; text-decoration:underline;'>Developed using RealVisXL_V5.0_Lightning model with ❤ by Aklavya</h2>"
+                "<h2 style='text-align:center; color:white; font-weight:bold; text-decoration:underline;'>Developed with ❤ by Aklavya</h2>"
             )
 
     with gr.Row():