Update app.py

app.py

@@ -518,7 +518,7 @@ def load_models() -> str | None:
                 index_path = _dl(MODEL_REPO, "transformer/diffusion_pytorch_model.safetensors.index.json")
                 with open(index_path) as handle:
                     index = _json.load(handle)
-                probe_key = "transformer_blocks.0.norm1.weight"
+                probe_key = "transformer_blocks.0.attn.to_q.weight"
                 shard_name = index["weight_map"][probe_key]
                 shard_path = _dl(MODEL_REPO, f"transformer/{shard_name}")
                 with safe_open(shard_path, framework="pt") as official_handle: