mirror of
https://github.com/Wan-Video/Wan2.1.git
synced 2025-11-04 06:15:17 +00:00
497 lines
20 KiB
Python
497 lines
20 KiB
Python
from __future__ import annotations
|
|
import os, io, tempfile, mimetypes
|
|
from typing import Any, Dict, List, Optional, Sequence, Tuple, Union, Literal
|
|
|
|
import gradio as gr
|
|
import PIL
|
|
from PIL import Image as PILImage
|
|
|
|
FilePath = str
|
|
ImageLike = Union["PIL.Image.Image", Any]
|
|
|
|
IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".bmp", ".gif", ".webp", ".tif", ".tiff", ".jfif", ".pjpeg"}
|
|
VIDEO_EXTS = {".mp4", ".mov", ".avi", ".mkv", ".webm", ".m4v", ".mpeg", ".mpg", ".ogv"}
|
|
|
|
def get_state(state):
|
|
return state if isinstance(state, dict) else state.value
|
|
|
|
def get_list( objs):
|
|
if objs is None:
|
|
return []
|
|
return [ obj[0] if isinstance(obj, tuple) else obj for obj in objs]
|
|
|
|
class AdvancedMediaGallery:
|
|
def __init__(
|
|
self,
|
|
label: str = "Media",
|
|
*,
|
|
media_mode: Literal["image", "video"] = "image",
|
|
height = None,
|
|
columns: Union[int, Tuple[int, ...]] = 6,
|
|
show_label: bool = True,
|
|
initial: Optional[Sequence[Union[FilePath, ImageLike]]] = None,
|
|
elem_id: Optional[str] = None,
|
|
elem_classes: Optional[Sequence[str]] = ("adv-media-gallery",),
|
|
accept_filter: bool = True, # restrict Add-button dialog to allowed extensions
|
|
single_image_mode: bool = False, # start in single-image mode (Add replaces)
|
|
):
|
|
assert media_mode in ("image", "video")
|
|
self.label = label
|
|
self.media_mode = media_mode
|
|
self.height = height
|
|
self.columns = columns
|
|
self.show_label = show_label
|
|
self.elem_id = elem_id
|
|
self.elem_classes = list(elem_classes) if elem_classes else None
|
|
self.accept_filter = accept_filter
|
|
|
|
items = self._normalize_initial(initial or [], media_mode)
|
|
|
|
# Components (filled on mount)
|
|
self.container: Optional[gr.Column] = None
|
|
self.gallery: Optional[gr.Gallery] = None
|
|
self.upload_btn: Optional[gr.UploadButton] = None
|
|
self.btn_remove: Optional[gr.Button] = None
|
|
self.btn_left: Optional[gr.Button] = None
|
|
self.btn_right: Optional[gr.Button] = None
|
|
self.btn_clear: Optional[gr.Button] = None
|
|
|
|
# Single dict state
|
|
self.state: Optional[gr.State] = None
|
|
self._initial_state: Dict[str, Any] = {
|
|
"items": items,
|
|
"selected": (len(items) - 1) if items else None,
|
|
"single": bool(single_image_mode),
|
|
"mode": self.media_mode,
|
|
}
|
|
|
|
# ---------------- helpers ----------------
|
|
|
|
def _normalize_initial(self, items: Sequence[Union[FilePath, ImageLike]], mode: str) -> List[Any]:
|
|
out: List[Any] = []
|
|
if mode == "image":
|
|
for it in items:
|
|
p = self._ensure_image_item(it)
|
|
if p is not None:
|
|
out.append(p)
|
|
else:
|
|
for it in items:
|
|
if isinstance(item, tuple): item = item[0]
|
|
if isinstance(it, str) and self._is_video_path(it):
|
|
out.append(os.path.abspath(it))
|
|
return out
|
|
|
|
def _ensure_image_item(self, item: Union[FilePath, ImageLike]) -> Optional[Any]:
|
|
# Accept a path to an image, or a PIL.Image/np.ndarray -> save temp PNG and return its path
|
|
if isinstance(item, tuple): item = item[0]
|
|
if isinstance(item, str):
|
|
return os.path.abspath(item) if self._is_image_path(item) else None
|
|
if PILImage is None:
|
|
return None
|
|
try:
|
|
if isinstance(item, PILImage.Image):
|
|
img = item
|
|
else:
|
|
import numpy as np # type: ignore
|
|
if isinstance(item, np.ndarray):
|
|
img = PILImage.fromarray(item)
|
|
elif hasattr(item, "read"):
|
|
data = item.read()
|
|
img = PILImage.open(io.BytesIO(data)).convert("RGBA")
|
|
else:
|
|
return None
|
|
tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
|
|
img.save(tmp.name)
|
|
return tmp.name
|
|
except Exception:
|
|
return None
|
|
|
|
@staticmethod
|
|
def _extract_path(obj: Any) -> Optional[str]:
|
|
# Try to get a filesystem path (for mode filtering); otherwise None.
|
|
if isinstance(obj, str):
|
|
return obj
|
|
try:
|
|
import pathlib
|
|
if isinstance(obj, pathlib.Path): # type: ignore
|
|
return str(obj)
|
|
except Exception:
|
|
pass
|
|
if isinstance(obj, dict):
|
|
return obj.get("path") or obj.get("name")
|
|
for attr in ("path", "name"):
|
|
if hasattr(obj, attr):
|
|
try:
|
|
val = getattr(obj, attr)
|
|
if isinstance(val, str):
|
|
return val
|
|
except Exception:
|
|
pass
|
|
return None
|
|
|
|
@staticmethod
|
|
def _is_image_path(p: str) -> bool:
|
|
ext = os.path.splitext(p)[1].lower()
|
|
if ext in IMAGE_EXTS:
|
|
return True
|
|
mt, _ = mimetypes.guess_type(p)
|
|
return bool(mt and mt.startswith("image/"))
|
|
|
|
@staticmethod
|
|
def _is_video_path(p: str) -> bool:
|
|
ext = os.path.splitext(p)[1].lower()
|
|
if ext in VIDEO_EXTS:
|
|
return True
|
|
mt, _ = mimetypes.guess_type(p)
|
|
return bool(mt and mt.startswith("video/"))
|
|
|
|
def _filter_items_by_mode(self, items: List[Any]) -> List[Any]:
|
|
# Enforce image-only or video-only collection regardless of how files were added.
|
|
out: List[Any] = []
|
|
if self.media_mode == "image":
|
|
for it in items:
|
|
p = self._extract_path(it)
|
|
if p is None:
|
|
# No path: likely an image object added programmatically => keep
|
|
out.append(it)
|
|
elif self._is_image_path(p):
|
|
out.append(os.path.abspath(p))
|
|
else:
|
|
for it in items:
|
|
p = self._extract_path(it)
|
|
if p is not None and self._is_video_path(p):
|
|
out.append(os.path.abspath(p))
|
|
return out
|
|
|
|
@staticmethod
|
|
def _concat_and_optionally_dedupe(cur: List[Any], add: List[Any]) -> List[Any]:
|
|
# Keep it simple: dedupe by path when available, else allow duplicates.
|
|
seen_paths = set()
|
|
def key(x: Any) -> Optional[str]:
|
|
if isinstance(x, str): return os.path.abspath(x)
|
|
try:
|
|
import pathlib
|
|
if isinstance(x, pathlib.Path): # type: ignore
|
|
return os.path.abspath(str(x))
|
|
except Exception:
|
|
pass
|
|
if isinstance(x, dict):
|
|
p = x.get("path") or x.get("name")
|
|
return os.path.abspath(p) if isinstance(p, str) else None
|
|
for attr in ("path", "name"):
|
|
if hasattr(x, attr):
|
|
try:
|
|
v = getattr(x, attr)
|
|
return os.path.abspath(v) if isinstance(v, str) else None
|
|
except Exception:
|
|
pass
|
|
return None
|
|
|
|
out: List[Any] = []
|
|
for lst in (cur, add):
|
|
for it in lst:
|
|
k = key(it)
|
|
if k is None or k not in seen_paths:
|
|
out.append(it)
|
|
if k is not None:
|
|
seen_paths.add(k)
|
|
return out
|
|
|
|
@staticmethod
|
|
def _paths_from_payload(payload: Any) -> List[Any]:
|
|
# Return as raw objects (paths/dicts/UploadedFile) to feed Gallery directly.
|
|
if payload is None:
|
|
return []
|
|
if isinstance(payload, (list, tuple, set)):
|
|
return list(payload)
|
|
return [payload]
|
|
|
|
# ---------------- event handlers ----------------
|
|
|
|
def _on_select(self, state: Dict[str, Any], gallery, evt: gr.SelectData) :
|
|
# Mirror the selected index into state and the gallery (server-side selected_index)
|
|
idx = None
|
|
if evt is not None and hasattr(evt, "index"):
|
|
ix = evt.index
|
|
if isinstance(ix, int):
|
|
idx = ix
|
|
elif isinstance(ix, (tuple, list)) and ix and isinstance(ix[0], int):
|
|
if isinstance(self.columns, int) and len(ix) >= 2:
|
|
idx = ix[0] * max(1, int(self.columns)) + ix[1]
|
|
else:
|
|
idx = ix[0]
|
|
st = get_state(state)
|
|
n = len(get_list(gallery))
|
|
sel = idx if (idx is not None and 0 <= idx < n) else None
|
|
st["selected"] = sel
|
|
return gr.update(selected_index=sel), st
|
|
|
|
def _on_gallery_change(self, value: List[Any], state: Dict[str, Any]) :
|
|
# Fires when users add/drag/drop/delete via the Gallery itself.
|
|
items_filtered = self._filter_items_by_mode(list(value or []))
|
|
st = get_state(state)
|
|
st["items"] = items_filtered
|
|
# Keep selection if still valid, else default to last
|
|
old_sel = st.get("selected", None)
|
|
if old_sel is None or not (0 <= old_sel < len(items_filtered)):
|
|
new_sel = (len(items_filtered) - 1) if items_filtered else None
|
|
else:
|
|
new_sel = old_sel
|
|
st["selected"] = new_sel
|
|
return gr.update(value=items_filtered, selected_index=new_sel), st
|
|
|
|
|
|
def _on_add(self, files_payload: Any, state: Dict[str, Any], gallery):
|
|
"""
|
|
Insert added items right AFTER the currently selected index.
|
|
Keeps the same ordering as chosen in the file picker, dedupes by path,
|
|
and re-selects the last inserted item.
|
|
"""
|
|
# New items (respect image/video mode)
|
|
new_items = self._filter_items_by_mode(self._paths_from_payload(files_payload))
|
|
|
|
st = get_state(state)
|
|
cur: List[Any] = get_list(gallery)
|
|
sel = st.get("selected", None)
|
|
if sel is None:
|
|
sel = (len(cur) -1) if len(cur)>0 else 0
|
|
single = bool(st.get("single", False))
|
|
|
|
# Nothing to add: keep as-is
|
|
if not new_items:
|
|
return gr.update(value=cur, selected_index=st.get("selected")), st
|
|
|
|
# Single-image mode: replace
|
|
if single:
|
|
st["items"] = [new_items[-1]]
|
|
st["selected"] = 0
|
|
return gr.update(value=st["items"], selected_index=0), st
|
|
|
|
# ---------- helpers ----------
|
|
def key_of(it: Any) -> Optional[str]:
|
|
# Prefer class helper if present
|
|
if hasattr(self, "_extract_path"):
|
|
p = self._extract_path(it) # type: ignore
|
|
else:
|
|
p = it if isinstance(it, str) else None
|
|
if p is None and isinstance(it, dict):
|
|
p = it.get("path") or it.get("name")
|
|
if p is None and hasattr(it, "path"):
|
|
try: p = getattr(it, "path")
|
|
except Exception: p = None
|
|
if p is None and hasattr(it, "name"):
|
|
try: p = getattr(it, "name")
|
|
except Exception: p = None
|
|
return os.path.abspath(p) if isinstance(p, str) else None
|
|
|
|
# Dedupe the incoming batch by path, preserve order
|
|
seen_new = set()
|
|
incoming: List[Any] = []
|
|
for it in new_items:
|
|
k = key_of(it)
|
|
if k is None or k not in seen_new:
|
|
incoming.append(it)
|
|
if k is not None:
|
|
seen_new.add(k)
|
|
|
|
# Remove any existing occurrences of the incoming items from current list,
|
|
# BUT keep the currently selected item even if it's also in incoming.
|
|
cur_clean: List[Any] = []
|
|
# sel_item = cur[sel] if (sel is not None and 0 <= sel < len(cur)) else None
|
|
# for idx, it in enumerate(cur):
|
|
# k = key_of(it)
|
|
# if it is sel_item:
|
|
# cur_clean.append(it)
|
|
# continue
|
|
# if k is not None and k in seen_new:
|
|
# continue # drop duplicate; we'll reinsert at the target spot
|
|
# cur_clean.append(it)
|
|
|
|
# # Compute insertion position: right AFTER the (possibly shifted) selected item
|
|
# if sel_item is not None:
|
|
# # find sel_item's new index in cur_clean
|
|
# try:
|
|
# pos_sel = cur_clean.index(sel_item)
|
|
# except ValueError:
|
|
# # Shouldn't happen, but fall back to end
|
|
# pos_sel = len(cur_clean) - 1
|
|
# insert_pos = pos_sel + 1
|
|
# else:
|
|
# insert_pos = len(cur_clean) # no selection -> append at end
|
|
insert_pos = min(sel, len(cur) -1)
|
|
cur_clean = cur
|
|
# Build final list and selection
|
|
merged = cur_clean[:insert_pos+1] + incoming + cur_clean[insert_pos+1:]
|
|
new_sel = insert_pos + len(incoming) # select the last inserted item
|
|
|
|
st["items"] = merged
|
|
st["selected"] = new_sel
|
|
return gr.update(value=merged, selected_index=new_sel), st
|
|
|
|
def _on_remove(self, state: Dict[str, Any], gallery) :
|
|
st = get_state(state); items: List[Any] = get_list(gallery); sel = st.get("selected", None)
|
|
if sel is None or not (0 <= sel < len(items)):
|
|
return gr.update(value=items, selected_index=st.get("selected")), st
|
|
items.pop(sel)
|
|
if not items:
|
|
st["items"] = []; st["selected"] = None
|
|
return gr.update(value=[], selected_index=None), st
|
|
new_sel = min(sel, len(items) - 1)
|
|
st["items"] = items; st["selected"] = new_sel
|
|
return gr.update(value=items, selected_index=new_sel), st
|
|
|
|
def _on_move(self, delta: int, state: Dict[str, Any], gallery) :
|
|
st = get_state(state); items: List[Any] = get_list(gallery); sel = st.get("selected", None)
|
|
if sel is None or not (0 <= sel < len(items)):
|
|
return gr.update(value=items, selected_index=sel), st
|
|
j = sel + delta
|
|
if j < 0 or j >= len(items):
|
|
return gr.update(value=items, selected_index=sel), st
|
|
items[sel], items[j] = items[j], items[sel]
|
|
st["items"] = items; st["selected"] = j
|
|
return gr.update(value=items, selected_index=j), st
|
|
|
|
def _on_clear(self, state: Dict[str, Any]) :
|
|
st = {"items": [], "selected": None, "single": state.get("single", False), "mode": self.media_mode}
|
|
return gr.update(value=[], selected_index=None), st
|
|
|
|
def _on_toggle_single(self, to_single: bool, state: Dict[str, Any]) :
|
|
st = get_state(state); st["single"] = bool(to_single)
|
|
items: List[Any] = list(st["items"]); sel = st.get("selected", None)
|
|
if st["single"]:
|
|
keep = items[sel] if (sel is not None and 0 <= sel < len(items)) else (items[-1] if items else None)
|
|
items = [keep] if keep is not None else []
|
|
sel = 0 if items else None
|
|
st["items"] = items; st["selected"] = sel
|
|
|
|
upload_update = gr.update(file_count=("single" if st["single"] else "multiple"))
|
|
left_update = gr.update(visible=not st["single"])
|
|
right_update = gr.update(visible=not st["single"])
|
|
clear_update = gr.update(visible=not st["single"])
|
|
gallery_update= gr.update(value=items, selected_index=sel)
|
|
|
|
return upload_update, left_update, right_update, clear_update, gallery_update, st
|
|
|
|
# ---------------- build & wire ----------------
|
|
|
|
def mount(self, parent: Optional[gr.Blocks | gr.Group | gr.Row | gr.Column] = None, update_form = False):
|
|
if parent is not None:
|
|
with parent:
|
|
col = self._build_ui()
|
|
else:
|
|
col = self._build_ui()
|
|
if not update_form:
|
|
self._wire_events()
|
|
return col
|
|
|
|
def _build_ui(self) -> gr.Column:
|
|
with gr.Column(elem_id=self.elem_id, elem_classes=self.elem_classes) as col:
|
|
self.container = col
|
|
|
|
self.state = gr.State(dict(self._initial_state))
|
|
|
|
self.gallery = gr.Gallery(
|
|
label=self.label,
|
|
value=self._initial_state["items"],
|
|
height=self.height,
|
|
columns=self.columns,
|
|
show_label=self.show_label,
|
|
preview= True,
|
|
selected_index=self._initial_state["selected"], # server-side selection
|
|
)
|
|
|
|
# One-line controls
|
|
exts = sorted(IMAGE_EXTS if self.media_mode == "image" else VIDEO_EXTS) if self.accept_filter else None
|
|
with gr.Row(equal_height=True, elem_classes=["amg-controls"]):
|
|
self.upload_btn = gr.UploadButton(
|
|
"Set" if self._initial_state["single"] else "Add",
|
|
file_types=exts,
|
|
file_count=("single" if self._initial_state["single"] else "multiple"),
|
|
variant="primary",
|
|
size="sm",
|
|
min_width=1,
|
|
)
|
|
self.btn_remove = gr.Button("Remove", size="sm", min_width=1)
|
|
self.btn_left = gr.Button("◀ Left", size="sm", visible=not self._initial_state["single"], min_width=1)
|
|
self.btn_right = gr.Button("Right ▶", size="sm", visible=not self._initial_state["single"], min_width=1)
|
|
self.btn_clear = gr.Button("Clear", variant="secondary", size="sm", visible=not self._initial_state["single"], min_width=1)
|
|
|
|
return col
|
|
|
|
def _wire_events(self):
|
|
# Selection: mirror into state and keep gallery.selected_index in sync
|
|
self.gallery.select(
|
|
self._on_select,
|
|
inputs=[self.state, self.gallery],
|
|
outputs=[self.gallery, self.state],
|
|
)
|
|
|
|
# Gallery value changed by user actions (click-to-add, drag-drop, internal remove, etc.)
|
|
self.gallery.change(
|
|
self._on_gallery_change,
|
|
inputs=[self.gallery, self.state],
|
|
outputs=[self.gallery, self.state],
|
|
)
|
|
|
|
# Add via UploadButton
|
|
self.upload_btn.upload(
|
|
self._on_add,
|
|
inputs=[self.upload_btn, self.state, self.gallery],
|
|
outputs=[self.gallery, self.state],
|
|
)
|
|
|
|
# Remove selected
|
|
self.btn_remove.click(
|
|
self._on_remove,
|
|
inputs=[self.state, self.gallery],
|
|
outputs=[self.gallery, self.state],
|
|
)
|
|
|
|
# Reorder using selected index, keep same item selected
|
|
self.btn_left.click(
|
|
lambda st, gallery: self._on_move(-1, st, gallery),
|
|
inputs=[self.state, self.gallery],
|
|
outputs=[self.gallery, self.state],
|
|
)
|
|
self.btn_right.click(
|
|
lambda st, gallery: self._on_move(+1, st, gallery),
|
|
inputs=[self.state, self.gallery],
|
|
outputs=[self.gallery, self.state],
|
|
)
|
|
|
|
# Clear all
|
|
self.btn_clear.click(
|
|
self._on_clear,
|
|
inputs=[self.state],
|
|
outputs=[self.gallery, self.state],
|
|
)
|
|
|
|
# ---------------- public API ----------------
|
|
|
|
def set_one_image_mode(self, enabled: bool = True):
|
|
"""Toggle single-image mode at runtime."""
|
|
return (
|
|
self._on_toggle_single,
|
|
[gr.State(enabled), self.state],
|
|
[self.upload_btn, self.btn_left, self.btn_right, self.btn_clear, self.gallery, self.state],
|
|
)
|
|
|
|
def get_toggable_elements(self):
|
|
return [self.upload_btn, self.btn_left, self.btn_right, self.btn_clear, self.gallery, self.state]
|
|
|
|
# import gradio as gr
|
|
|
|
# with gr.Blocks() as demo:
|
|
# amg = AdvancedMediaGallery(media_mode="image", height=190, columns=8)
|
|
# amg.mount()
|
|
# g = amg.gallery
|
|
# # buttons to switch modes live (optional)
|
|
# def process(g):
|
|
# pass
|
|
# with gr.Row():
|
|
# gr.Button("toto").click(process, g)
|
|
# gr.Button("ONE image").click(*amg.set_one_image_mode(True))
|
|
# gr.Button("MULTI image").click(*amg.set_one_image_mode(False))
|
|
|
|
# demo.launch()
|