SEN-10: propagate image patch from fhdr-uncensored (fallback anti-gated models)

app.py

@@ -17,9 +17,20 @@ from diffusers.utils import export_to_video
 SPACE_ID = os.getenv("SPACE_ID", "").lower()
 IS_VIDEO_SPACE = any(k in SPACE_ID for k in ["hunyuanvideo", "wan-2-1"])
 
-IMAGE_MODEL_ID = os.getenv("IMAGE_MODEL_ID", "black-forest-labs/FLUX.1-schnell")
+IMAGE_MODEL_ID = os.getenv("IMAGE_MODEL_ID", "runwayml/stable-diffusion-v1-5")
 VIDEO_MODEL_ID = os.getenv("VIDEO_MODEL_ID", "damo-vilab/text-to-video-ms-1.7b")
 
+# Known ungated defaults per space: avoids GatedRepoError on HF Spaces without manual model-license acceptance.
+SPACE_IMAGE_DEFAULTS = {
+    "fhdr-uncensored": "SG161222/Realistic_Vision_V6.0_B1_noVAE",
+    "z-image-turbo": "stabilityai/sdxl-turbo",
+}
+FALLBACK_IMAGE_MODELS = [
+    IMAGE_MODEL_ID,
+    SPACE_IMAGE_DEFAULTS.get(SPACE_ID.split("/")[-1], ""),
+    "runwayml/stable-diffusion-v1-5",
+]
+
 _image_pipe: Optional[DiffusionPipeline] = None
 _video_pipe: Optional[DiffusionPipeline] = None
 
@@ -36,7 +47,33 @@ def _load_image_pipe() -> DiffusionPipeline:
     global _image_pipe
     if _image_pipe is None:
         device, dtype = _device_dtype()
-        _image_pipe = DiffusionPipeline.from_pretrained(IMAGE_MODEL_ID, torch_dtype=dtype)
+        last_error = None
+        for model_id in [m for m in FALLBACK_IMAGE_MODELS if m]:
+            try:
+                # Try to disable safety checker when supported.
+                _image_pipe = DiffusionPipeline.from_pretrained(
+                    model_id,
+                    torch_dtype=dtype,
+                    safety_checker=None,
+                    requires_safety_checker=False,
+                )
+                break
+            except TypeError:
+                _image_pipe = DiffusionPipeline.from_pretrained(model_id, torch_dtype=dtype)
+                break
+            except Exception as e:
+                last_error = e
+                _image_pipe = None
+
+        if _image_pipe is None:
+            raise RuntimeError(f"Unable to load image model from {FALLBACK_IMAGE_MODELS}: {last_error}")
+
+        # Explicit runtime bypass for diffusion pipelines exposing an NSFW safety checker.
+        if hasattr(_image_pipe, "safety_checker"):
+            _image_pipe.safety_checker = None
+        if hasattr(_image_pipe, "requires_safety_checker"):
+            _image_pipe.requires_safety_checker = False
+
         if device == "cuda":
             _image_pipe.enable_model_cpu_offload()
         else:
@@ -73,13 +110,17 @@ def generate_image(prompt: str, steps: int, guidance_scale: float, seed: int):
     pipe = _load_image_pipe()
 
     gen = torch.Generator(device="cpu").manual_seed(int(seed))
+    default_size = 512 if "v1-5" in str(getattr(pipe, "name_or_path", "")).lower() else 1024
+    width = int(os.getenv("IMAGE_WIDTH", default_size))
+    height = int(os.getenv("IMAGE_HEIGHT", default_size))
+
     image: Image.Image = pipe(
         prompt=prompt,
         num_inference_steps=int(steps),
         guidance_scale=float(guidance_scale),
         generator=gen,
-        width=1024,
-        height=1024,
+        width=width,
+        height=height,
     ).images[0]
     return image