Update app.py

app.py

@@ -227,6 +227,7 @@ 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")
@@ -234,25 +235,35 @@ _sd_t2i = None
 _sd_i2i = None
 
 def _lazy_sd_pipes():
-    """Load SD once without low_cpu_mem_usage to avoid offload_state_dict kwarg; reuse modules for img2img."""
+    """
+    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)
+    """
     global _sd_t2i, _sd_i2i
     if _sd_t2i is not None and _sd_i2i is not None:
         return _sd_t2i, _sd_i2i
 
-    dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+    use_cuda = torch.cuda.is_available()
+    dtype = torch.float16 if use_cuda else torch.float32
 
+    # Important: torch_dtype + low_cpu_mem_usage=False to avoid offload_state_dict
     _sd_t2i = StableDiffusionPipeline.from_pretrained(
         SD_MODEL,
-        dtype=dtype,                 # (`torch_dtype` is deprecated, use `dtype`)
-        safety_checker=None,
+        torch_dtype=dtype,            # <-- correct kwarg
+        safety_checker=None,          # disable safety checker at load
         feature_extractor=None,
         use_safetensors=True,
-        low_cpu_mem_usage=False      # <-- critical: prevents passing offload_state_dict
+        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
     )
-    if torch.cuda.is_available():
+
+    if use_cuda:
         _sd_t2i = _sd_t2i.to("cuda")
 
-    # Build img2img from already-loaded modules (avoids another from_pretrained call)
+    # Build img2img from the same components (avoid another from_pretrained call)
     _sd_i2i = StableDiffusionImg2ImgPipeline(
         vae=_sd_t2i.vae,
         text_encoder=_sd_t2i.text_encoder,
@@ -262,12 +273,11 @@ def _lazy_sd_pipes():
         safety_checker=None,
         feature_extractor=None
     )
-    if torch.cuda.is_available():
+    if use_cuda:
         _sd_i2i = _sd_i2i.to("cuda")
 
     return _sd_t2i, _sd_i2i
 
-
 def _save_keyframe(pid: str, shot_id: int, img: Image.Image) -> str:
     pdir = project_dir(pid)
     out = os.path.join(pdir, "keyframes", f"shot_{shot_id:02d}.png")