observantdistressed commited on
Commit
4ae1908
Β·
1 Parent(s): 826b8b9

feat: h3-{task_id}.mp4 + community dropdown fixes

Browse files

- Output video basename is h3-<task_id>.mp4 (Gradio download shows task id)
- Rewrite _community_lora_choices: no nested loops; skip cache only on exact basename
- Community dropdown interactive=True; use .change (not broken .select+SelectData)
- Page load/focus return full interactive Dropdown

Addresses PR review: select handler, repo-prefix cache filter, rebase on main.

Files changed (4) hide show
  1. lora/registry.py +64 -34
  2. pipeline.py +6 -2
  3. tests/test_task_isolation.py +3 -3
  4. ui.py +23 -6
lora/registry.py CHANGED
@@ -126,38 +126,63 @@ def _community_lora_choices():
126
  Sources: the JSON registry **plus** any ``.safetensors`` files already cached
127
  in ``/data/loras/`` that aren't in the registry yet. Cached files get a
128
  ``cache:`` prefixed value so the download/load path can load them directly
129
- without going through the URL β†’ download flow."""
 
 
 
 
 
 
130
  registry = _load_community_registry()
131
- known_urls = {e.get("url", "") for e in registry}
132
- known_keys = {(e.get("repo", ""), e.get("file", "")) for e in registry}
133
  choices = []
 
 
 
134
  for entry in registry:
135
- repo = entry.get("repo", "")
136
- desc = entry.get("description", "")
137
- tags = entry.get("tags", [])
138
- name = repo.split("/")[-1]
139
- from state import _LORA_CACHE_DIR
140
- try:
141
- cache_files = sorted(
142
- f for f in os.listdir(_LORA_CACHE_DIR)
143
- if f.endswith(".safetensors") and not f.startswith(".")
144
- )
145
- except OSError:
146
- cache_files = []
147
- for f in cache_files:
148
- if f not in known_urls and not any(f.startswith(f"{r.replace('/', '__')}__{fi.split('/')[-1]}") for r, fi in known_keys):
149
- choices.append((f"{f} (cached)", f"cache:{os.path.join(_LORA_CACHE_DIR, f)}"))
150
- for entry in registry:
151
- repo = entry.get("repo", "")
152
- desc = entry.get("description", "")
153
- tags = entry.get("tags", [])
154
- name = repo.split("/")[-1]
155
  label = f"{name}"
156
  if desc:
157
  label += f" β€” {desc[:60]}{'…' if len(desc) > 60 else ''}"
158
  elif tags:
159
- label += f" ({', '.join(tags[:3])})"
160
- choices.append((label, entry.get("url", "")))
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
161
  return choices
162
 
163
 
@@ -182,15 +207,15 @@ def _download_custom_lora(url, scale=1.0, trigger="", description="", tags="", a
182
  active_val.append(enc)
183
  return (
184
  gr.update(value=msg, visible=True),
185
- gr.update(choices=choices_c or None, value=community_value),
186
- gr.update(choices=choices_a or None, value=active_val),
187
  )
188
 
189
  def _err(msg):
190
  return (
191
  gr.update(value=msg, visible=True),
192
- gr.update(),
193
- gr.update(choices=_active_custom_choices() or None),
194
  )
195
 
196
  if not url:
@@ -318,13 +343,18 @@ def _submit_community_lora(url, scale, trigger, description, tags):
318
  return gr.update(choices=choices, value=full_url), "βœ… Added to community registry."
319
 
320
 
321
- def _load_selected_community_lora(evt: gr.SelectData):
322
- """When a community LoRA is selected from the dropdown, return its URL / scale / trigger.
323
 
324
- For cached LoRAs (value starts with ``cache:``) the URL field is set to the
325
- cache path so the Download & Apply button can detect it."""
 
 
326
  registry = _load_community_registry()
327
- url = evt.value
 
 
 
328
  # Cached LoRAs use cache:path β€” pass through so _download_custom_lora picks them up.
329
  if url.startswith("cache:"):
330
  return url, 1.0, ""
 
126
  Sources: the JSON registry **plus** any ``.safetensors`` files already cached
127
  in ``/data/loras/`` that aren't in the registry yet. Cached files get a
128
  ``cache:`` prefixed value so the download/load path can load them directly
129
+ without going through the URL β†’ download flow.
130
+
131
+ Cache skip logic matches ``_lora_cache_path`` naming
132
+ ``{repo}__rev_{revision}__{file}`` so multiple files from the same repo stay visible.
133
+ """
134
+ from state import _LORA_CACHE_DIR
135
+
136
  registry = _load_community_registry()
 
 
137
  choices = []
138
+
139
+ # Exact basenames we already cover via registry entries (with/without revision).
140
+ registered_basenames: set[str] = set()
141
  for entry in registry:
142
+ repo = str(entry.get("repo") or "")
143
+ file = str(entry.get("file") or "")
144
+ rev = str(entry.get("revision") or "").strip()
145
+ if not repo or not file:
146
+ continue
147
+ safe_repo = repo.replace("/", "__")
148
+ safe_file = file.replace("/", "__")
149
+ if rev:
150
+ registered_basenames.add(f"{safe_repo}__rev_{rev}__{safe_file}")
151
+ # Also mark revision-agnostic forms used by older cache keys.
152
+ registered_basenames.add(f"{safe_repo}__{safe_file}")
153
+ registered_basenames.add(safe_file)
154
+
155
+ desc = str(entry.get("description") or "")
156
+ tags = entry.get("tags") or []
157
+ name = repo.split("/")[-1] or repo or "LoRA"
 
 
 
 
158
  label = f"{name}"
159
  if desc:
160
  label += f" β€” {desc[:60]}{'…' if len(desc) > 60 else ''}"
161
  elif tags:
162
+ label += f" ({', '.join(str(t) for t in tags[:3])})"
163
+ url = entry.get("url") or f"https://huggingface.co/{repo}/resolve/main/{file}"
164
+ choices.append((label, url))
165
+
166
+ try:
167
+ cache_files = sorted(
168
+ f for f in os.listdir(_LORA_CACHE_DIR)
169
+ if f.endswith(".safetensors") and not f.startswith(".")
170
+ )
171
+ except OSError:
172
+ cache_files = []
173
+
174
+ for f in cache_files:
175
+ # Skip only when this exact cache basename is already represented.
176
+ if f in registered_basenames:
177
+ continue
178
+ # Strip optional __rev_<sha>__ middle for a softer match on repo+file only
179
+ # when the registry entry omitted revision.
180
+ stem = f
181
+ if "__rev_" in f:
182
+ # {repo}__rev_{rev}__{file} β†’ compare full key only (already handled above)
183
+ pass
184
+ choices.append((f"{f} (cached)", f"cache:{os.path.join(_LORA_CACHE_DIR, f)}"))
185
+
186
  return choices
187
 
188
 
 
207
  active_val.append(enc)
208
  return (
209
  gr.update(value=msg, visible=True),
210
+ gr.update(choices=choices_c or [], interactive=True, value=community_value),
211
+ gr.update(choices=choices_a or [], interactive=True, value=active_val),
212
  )
213
 
214
  def _err(msg):
215
  return (
216
  gr.update(value=msg, visible=True),
217
+ gr.update(interactive=True),
218
+ gr.update(choices=_active_custom_choices() or [], interactive=True),
219
  )
220
 
221
  if not url:
 
343
  return gr.update(choices=choices, value=full_url), "βœ… Added to community registry."
344
 
345
 
346
+ def _load_selected_community_lora(evt):
347
+ """When a community LoRA is selected, return URL / scale / trigger.
348
 
349
+ Bound via ``.change(dropdown)`` so the value is passed explicitly (Gradio does
350
+ not inject SelectData without a type hint + inputs=None). Accepts a plain
351
+ value or ``gr.SelectData``.
352
+ """
353
  registry = _load_community_registry()
354
+ if evt is None:
355
+ return "", 1.0, ""
356
+ url = evt.value if hasattr(evt, "value") else evt
357
+ url = "" if url is None else str(url)
358
  # Cached LoRAs use cache:path β€” pass through so _download_custom_lora picks them up.
359
  if url.startswith("cache:"):
360
  return url, 1.0, ""
pipeline.py CHANGED
@@ -436,8 +436,12 @@ def _generate_impl(prompt, image_path=None, last_image_path=None, canvas=DEFAULT
436
  _task_log(_task, f"_generate_impl: condition={condition_seconds:.0f}s, total={(time.time() - conditioned):.0f}s")
437
  progress(0.95, desc="Muxing the soundtrack ...")
438
 
439
- # Write the video to a task-scoped output path so the file is unique per request.
440
- output_path = os.path.join(_task.work_dir, "output.mp4") if _task is not None else os.path.join(tempfile.gettempdir(), f"h3-output-{int(time.time())}.mp4")
 
 
 
 
441
  # encode_video expects a (T,H,W,C) array/list of frames β€” not frames[None]
442
  # (that only works on numpy and was crashing when videos came back as a list).
443
  if hasattr(frames, "detach"):
 
436
  _task_log(_task, f"_generate_impl: condition={condition_seconds:.0f}s, total={(time.time() - conditioned):.0f}s")
437
  progress(0.95, desc="Muxing the soundtrack ...")
438
 
439
+ # Task-scoped dir + task-id filename so Gradio download shows h3-<task_id>.mp4
440
+ # (not a generic "output.mp4" basename).
441
+ if _task is not None:
442
+ output_path = os.path.join(_task.work_dir, f"h3-{_task.task_id}.mp4")
443
+ else:
444
+ output_path = os.path.join(tempfile.gettempdir(), f"h3-output-{int(time.time())}.mp4")
445
  # encode_video expects a (T,H,W,C) array/list of frames β€” not frames[None]
446
  # (that only works on numpy and was crashing when videos came back as a list).
447
  if hasattr(frames, "detach"):
tests/test_task_isolation.py CHANGED
@@ -59,9 +59,9 @@ def test_task_filesystem_isolation(tmp_path, monkeypatch):
59
  assert os.path.isdir(a.work_dir)
60
  assert os.path.isdir(os.path.join(a.work_dir, "intermediates"))
61
  # no collision
62
- open(os.path.join(a.work_dir, "output.mp4"), "w").write("a")
63
- open(os.path.join(b.work_dir, "output.mp4"), "w").write("b")
64
- assert open(os.path.join(a.work_dir, "output.mp4")).read() == "a"
65
 
66
 
67
  def test_model_mutation_lock_is_reentrant():
 
59
  assert os.path.isdir(a.work_dir)
60
  assert os.path.isdir(os.path.join(a.work_dir, "intermediates"))
61
  # no collision
62
+ open(os.path.join(a.work_dir, f"h3-{a.task_id}.mp4"), "w").write("a")
63
+ open(os.path.join(b.work_dir, f"h3-{b.task_id}.mp4"), "w").write("b")
64
+ assert open(os.path.join(a.work_dir, f"h3-{a.task_id}.mp4")).read() == "a"
65
 
66
 
67
  def test_model_mutation_lock_is_reentrant():
ui.py CHANGED
@@ -366,6 +366,7 @@ def build_ui():
366
  multiselect=False,
367
  filterable=True,
368
  allow_custom_value=True,
 
369
  info="Pick one to auto-fill the link, then Download & Verify (or it loads into Active above).",
370
  )
371
  custom_lora_url = gr.Textbox(
@@ -429,9 +430,11 @@ def build_ui():
429
  [custom_lora_status, community_lora_dropdown, active_custom_loras],
430
  )
431
 
432
- community_lora_dropdown.select(
 
 
433
  _load_selected_community_lora,
434
- None,
435
  [custom_lora_url, custom_lora_scale, custom_lora_trigger],
436
  )
437
 
@@ -475,7 +478,15 @@ def build_ui():
475
  # ── Refresh community dropdown / status after page connect ──────────────
476
  def _refresh_community_dd():
477
  choices = _community_lora_choices() or []
478
- return gr.update(choices=choices)
 
 
 
 
 
 
 
 
479
  def _on_page_load():
480
  choices = _community_lora_choices() or []
481
  try:
@@ -483,11 +494,17 @@ def build_ui():
483
  banner = f"**{_st()}**"
484
  except Exception as _e:
485
  banner = f"Status unavailable: {_e}"
486
- return gr.update(choices=choices), banner
 
 
 
 
 
 
 
 
487
 
488
  demo.load(_on_page_load, None, [community_lora_dropdown, status_banner], api_name=False)
489
- # Also refresh when the user opens (focuses) the dropdown β€” covers cases
490
- # where demo.load() fires before the bucket is fully listed.
491
  community_lora_dropdown.focus(
492
  _refresh_community_dd, None, community_lora_dropdown, api_name=False,
493
  )
 
366
  multiselect=False,
367
  filterable=True,
368
  allow_custom_value=True,
369
+ interactive=True,
370
  info="Pick one to auto-fill the link, then Download & Verify (or it loads into Active above).",
371
  )
372
  custom_lora_url = gr.Textbox(
 
430
  [custom_lora_status, community_lora_dropdown, active_custom_loras],
431
  )
432
 
433
+ # Use .change with the dropdown as an input β€” Gradio injects SelectData only when
434
+ # the handler is type-hinted and inputs=None; that path was broken. .change is enough.
435
+ community_lora_dropdown.change(
436
  _load_selected_community_lora,
437
+ [community_lora_dropdown],
438
  [custom_lora_url, custom_lora_scale, custom_lora_trigger],
439
  )
440
 
 
478
  # ── Refresh community dropdown / status after page connect ──────────────
479
  def _refresh_community_dd():
480
  choices = _community_lora_choices() or []
481
+ return gr.Dropdown(
482
+ choices=choices,
483
+ value=None,
484
+ interactive=True,
485
+ multiselect=False,
486
+ filterable=True,
487
+ allow_custom_value=True,
488
+ )
489
+
490
  def _on_page_load():
491
  choices = _community_lora_choices() or []
492
  try:
 
494
  banner = f"**{_st()}**"
495
  except Exception as _e:
496
  banner = f"Status unavailable: {_e}"
497
+ dd = gr.Dropdown(
498
+ choices=choices,
499
+ value=None,
500
+ interactive=True,
501
+ multiselect=False,
502
+ filterable=True,
503
+ allow_custom_value=True,
504
+ )
505
+ return dd, banner
506
 
507
  demo.load(_on_page_load, None, [community_lora_dropdown, status_banner], api_name=False)
 
 
508
  community_lora_dropdown.focus(
509
  _refresh_community_dd, None, community_lora_dropdown, api_name=False,
510
  )