Update app.py

app.py

@@ -227,20 +227,20 @@ def generate_storyboard_with_llm(user_prompt: str, n_shots: int, default_fps: in
 # =========================
 # IMAGE GEN (ZeroGPU) — SD1.5 text2img + img2img chaining
 # =========================
-# --- replace the whole _lazy_sd_pipes() with this ---
+
 from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline
 
-SD_MODEL = os.getenv("SD_MODEL", "runwayml/stable-diffusion-v1-5")
+SD_MODEL = os.getenv("SD_MODEL", "stabilityai/sd-turbo")
 _sd_t2i = None
 _sd_i2i = None
 
 def _lazy_sd_pipes():
     """
-    Load SD once with settings that avoid the `offload_state_dict` kwarg issue:
-      - use `torch_dtype` (not `dtype`)
-      - set `low_cpu_mem_usage=False`
-      - disable safety checker at load time
-      - build Img2Img from the already-loaded modules (no second from_pretrained)
+    Load SD once in a version-safe way:
+      - torch_dtype (not dtype)
+      - low_cpu_mem_usage=False to avoid offload_state_dict kwarg
+      - no revision pin (some repos don't have 'fp16' branch)
+      - optional HF token if set (for gated models)
     """
     global _sd_t2i, _sd_i2i
     if _sd_t2i is not None and _sd_i2i is not None:
@@ -248,22 +248,20 @@ def _lazy_sd_pipes():
 
     use_cuda = torch.cuda.is_available()
     dtype = torch.float16 if use_cuda else torch.float32
+    hf_token = os.getenv("HF_TOKEN", None)  # add this in Space Secrets only if needed
 
-    # Important: torch_dtype + low_cpu_mem_usage=False to avoid offload_state_dict
     _sd_t2i = StableDiffusionPipeline.from_pretrained(
         SD_MODEL,
-        torch_dtype=dtype,            # <-- correct kwarg
-        safety_checker=None,          # disable safety checker at load
+        torch_dtype=dtype,
+        safety_checker=None,
         feature_extractor=None,
         use_safetensors=True,
-        low_cpu_mem_usage=False,      # <-- critical to prevent passing offload_state_dict
-        revision="fp16" if use_cuda else None  # helps pick fp16 weights on GPU
+        low_cpu_mem_usage=False,
+        token=hf_token
     )
-
     if use_cuda:
         _sd_t2i = _sd_t2i.to("cuda")
 
-    # Build img2img from the same components (avoid another from_pretrained call)
     _sd_i2i = StableDiffusionImg2ImgPipeline(
         vae=_sd_t2i.vae,
         text_encoder=_sd_t2i.text_encoder,