import os # PYTORCH_CUDA_ALLOC_CONF MUSS vor torch-Import stehen (GPT-Fix) os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True" import spaces import sys import torch import shutil import tempfile import gradio as gr from PIL import Image from rembg import remove import uuid import subprocess from glob import glob from huggingface_hub import snapshot_download import zipfile import json from pathlib import Path def _get_subprocess_env() -> dict: """Subprocess-Umgebung mit CUDA-Library-Pfaden und NVML-freiem Allocator. backend:cudaMallocAsync ersetzt CUDACachingAllocator komplett – kein NVML nötig, stabil auf ZeroGPU Blackwell + torch 2.11.0 + CUDA 13. """ import site import glob as _glob env = os.environ.copy() # CUDA-Library-Pfade für libcudart nvidia_paths: list[str] = [] all_site = site.getsitepackages() try: all_site = all_site + [site.getusersitepackages()] except Exception: pass for sp in all_site: matches = _glob.glob(os.path.join(sp, "nvidia", "*", "lib")) nvidia_paths.extend(m for m in matches if os.path.isdir(m)) torch_lib = os.path.join(os.path.dirname(torch.__file__), "lib") system_cuda = [ "/usr/local/cuda-13/lib64", "/usr/local/cuda-13.0/lib64", "/usr/local/cuda-12/lib64", "/usr/local/cuda/lib64", "/usr/lib/x86_64-linux-gnu", ] extra = [p for p in system_cuda if os.path.isdir(p)] existing_ld = env.get("LD_LIBRARY_PATH", "") env["LD_LIBRARY_PATH"] = ":".join( nvidia_paths + [torch_lib] + extra + ([existing_ld] if existing_ld else []) ) # NVML-freier Allocator für Subprocess (CUDACachingAllocator-Assertions vermeiden) env["PYTORCH_CUDA_ALLOC_CONF"] = "backend:cudaMallocAsync" return env def _validate_xformers() -> bool: """Prüft ob xformers.ops.memory_efficient_attention wirklich funktioniert. Gibt True zurück wenn ein Mini-Test mit echten CUDA-Tensoren besteht. """ try: import xformers import xformers.ops print(f"[DIAG] xformers version : {xformers.__version__}", flush=True) has_mea = hasattr(xformers.ops, 'memory_efficient_attention') print(f"[DIAG] xformers MEA attr : {has_mea}", flush=True) if not has_mea: print("[DIAG] xformers MEA test : SKIP (attr fehlt)", flush=True) return False if not torch.cuda.is_available(): print("[DIAG] xformers MEA test : SKIP (kein CUDA)", flush=True) return False # Mini-Test mit kleinen CUDA-Tensoren q = torch.randn(2, 16, 64, device="cuda", dtype=torch.float16) k = torch.randn(2, 16, 64, device="cuda", dtype=torch.float16) v = torch.randn(2, 16, 64, device="cuda", dtype=torch.float16) _ = xformers.ops.memory_efficient_attention(q, k, v) torch.cuda.synchronize() print("[DIAG] xformers MEA test : PASS ✅", flush=True) return True except Exception as e: print(f"[DIAG] xformers MEA test : FAIL – {e}", flush=True) return False def _log_cuda_diagnostics(): import glob as _glob import diffusers print(f"[DIAG] diffusers version : {diffusers.__version__}", flush=True) print(f"[DIAG] torch version : {torch.__version__}", flush=True) print(f"[DIAG] torch CUDA build : {torch.version.cuda}", flush=True) print(f"[DIAG] CUDA available : {torch.cuda.is_available()}", flush=True) print(f"[DIAG] sys.executable : {sys.executable}", flush=True) env = _get_subprocess_env() ld = env.get("LD_LIBRARY_PATH", "") found = [h for d in ld.split(":") if d for h in _glob.glob(os.path.join(d, "libcudart.so*"))] print(f"[DIAG] libcudart found : {found}", flush=True) print(f"[DIAG] subprocess ALLOC : {env.get('PYTORCH_CUDA_ALLOC_CONF')}", flush=True) _validate_xformers() import importlib.util print(f"[DIAG] pytorch3d spec : {importlib.util.find_spec('pytorch3d')}", flush=True) _log_cuda_diagnostics() os.makedirs("ckpts", exist_ok=True) snapshot_download(repo_id="pengHTYX/PSHuman_Unclip_768_6views", local_dir="./ckpts") os.makedirs("smpl_related", exist_ok=True) snapshot_download(repo_id="fffiloni/PSHuman-SMPL-related", local_dir="./smpl_related") examples_folder = "examples" images_examples = [ os.path.join(examples_folder, file) for file in os.listdir(examples_folder) if os.path.isfile(os.path.join(examples_folder, file)) ] ALLOWED_IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".webp"} def get_uploaded_file_path(file_obj) -> str: """Return a filesystem path from Gradio File outputs across Gradio versions.""" if file_obj is None: raise gr.Error("Bitte zuerst eine ZIP-Datei auswählen.") if isinstance(file_obj, (str, os.PathLike)): return str(file_obj) if isinstance(file_obj, dict): path = file_obj.get("path") or file_obj.get("name") if path: return str(path) path = getattr(file_obj, "name", None) or getattr(file_obj, "path", None) if path: return str(path) raise gr.Error("ZIP-Dateipfad konnte nicht gelesen werden.") def remove_background(input_pil, remove_bg): temp_dir = tempfile.mkdtemp(prefix="pshuman_session_") unique_id = str(uuid.uuid4()) image_path = os.path.join(temp_dir, f"input_image_{unique_id}.png") try: if isinstance(input_pil, Image.Image): image = input_pil else: image = Image.open(input_pil) image = image.transpose(Image.FLIP_LEFT_RIGHT) image.save(image_path) except Exception as e: shutil.rmtree(temp_dir, ignore_errors=True) raise gr.Error(f"Fehler beim Laden/Speichern des Bildes: {str(e)}") if remove_bg is True: removed_bg_path = os.path.join(temp_dir, f"output_image_rmbg_{unique_id}.png") try: img = Image.open(image_path) result = remove(img) result.save(removed_bg_path) os.remove(image_path) except Exception as e: shutil.rmtree(temp_dir, ignore_errors=True) raise gr.Error(f"Fehler bei der Hintergrundentfernung: {str(e)}") return removed_bg_path, temp_dir else: return image_path, temp_dir def get_multiview_root(session_dir: str) -> Path: return Path(session_dir) / "multiview" def find_single_scene_dir(session_dir: str) -> Path: mv_root = get_multiview_root(session_dir) if not mv_root.exists(): raise gr.Error(f"Kein multiview-Ordner gefunden: {mv_root}") scene_dirs = [p for p in mv_root.iterdir() if p.is_dir()] if not scene_dirs: raise gr.Error(f"Keine Szene im multiview-Ordner gefunden: {mv_root}") if len(scene_dirs) > 1: raise gr.Error("Mehrere Szenen gefunden. Diese App erwartet aktuell genau eine Szene pro Session.") return scene_dirs[0] def get_scene_name_from_meta(meta_path: Path) -> str: if not meta_path.exists(): return "uploaded_scene" try: with open(meta_path, "r", encoding="utf-8") as f: data = json.load(f) scene_name = str(data.get("scene", "uploaded_scene")).strip() return scene_name or "uploaded_scene" except Exception: return "uploaded_scene" def ensure_session_from_upload(session_dir: str | None, zip_file_path: str) -> str: if session_dir and Path(session_dir).exists(): return session_dir new_session_dir = tempfile.mkdtemp(prefix="pshuman_session_") scene_name = "uploaded_scene" try: with zipfile.ZipFile(zip_file_path, "r") as zf: if "meta.json" in zf.namelist(): tmp_meta_dir = Path(new_session_dir) / "_meta_tmp" tmp_meta_dir.mkdir(parents=True, exist_ok=True) zf.extract("meta.json", path=tmp_meta_dir) scene_name = get_scene_name_from_meta(tmp_meta_dir / "meta.json") shutil.rmtree(tmp_meta_dir, ignore_errors=True) except Exception: scene_name = "uploaded_scene" scene_dir = Path(new_session_dir) / "multiview" / scene_name (scene_dir / "edit").mkdir(parents=True, exist_ok=True) (scene_dir / "raw").mkdir(parents=True, exist_ok=True) meta_path = scene_dir / "meta.json" if not meta_path.exists(): with open(meta_path, "w", encoding="utf-8") as f: json.dump({"scene": scene_name, "source": "upload_created_session"}, f, indent=2) return new_session_dir def list_gallery_images(session_dir: str): if not session_dir or not Path(session_dir).exists(): return [] mv_root = get_multiview_root(session_dir) if not mv_root.exists(): return [] try: scene_dir = find_single_scene_dir(session_dir) edit_dir = scene_dir / "edit" if not edit_dir.exists(): return [] return sorted(str(p) for p in edit_dir.glob("color_*") if p.is_file()) except Exception: return [] def create_edit_zip(session_dir: str) -> str: scene_dir = find_single_scene_dir(session_dir) edit_dir = scene_dir / "edit" meta_path = scene_dir / "meta.json" if not edit_dir.exists(): raise gr.Error(f"Kein edit-Ordner gefunden: {edit_dir}") zip_path = Path(session_dir) / f"{scene_dir.name}_multiview_edit.zip" with zipfile.ZipFile(zip_path, "w", compression=zipfile.ZIP_DEFLATED) as zf: for file_path in sorted(edit_dir.iterdir()): if file_path.is_file(): zf.write(file_path, arcname=file_path.name) if meta_path.exists(): zf.write(meta_path, arcname="meta.json") return str(zip_path) def inspect_edit_set(edit_dir: Path) -> dict: existing_files = {p.name for p in edit_dir.iterdir() if p.is_file()} color_files = sorted([ f for f in existing_files if f.startswith("color_") and Path(f).suffix.lower() in ALLOWED_IMAGE_EXTENSIONS ]) normal_files = sorted([ f for f in existing_files if f.startswith("normal_") and Path(f).suffix.lower() in ALLOWED_IMAGE_EXTENSIONS ]) if color_files: expected_indices = sorted([f.split("_")[1].split(".")[0] for f in color_files]) elif normal_files: expected_indices = sorted([f.split("_")[1].split(".")[0] for f in normal_files]) else: expected_indices = [] expected_color_names = {f"color_{idx}.png" for idx in expected_indices} expected_normal_names = {f"normal_{idx}.png" for idx in expected_indices} found_colors = {f for f in existing_files if f.startswith("color_")} found_normals = {f for f in existing_files if f.startswith("normal_")} missing_colors = sorted(expected_color_names - found_colors) if expected_indices else [] missing_normals = sorted(expected_normal_names - found_normals) if expected_indices else [] has_full_colors = bool(expected_indices) and len(missing_colors) == 0 has_full_normals = bool(expected_indices) and len(missing_normals) == 0 if has_full_colors and has_full_normals: state = "READY" elif has_full_colors and not has_full_normals: state = "NEEDS_NORMALS" else: state = "INVALID" return { "state": state, "expected_indices": expected_indices, "found_colors": sorted(found_colors), "found_normals": sorted(found_normals), "missing_colors": missing_colors, "missing_normals": missing_normals, "has_full_colors": has_full_colors, "has_full_normals": has_full_normals, } def validate_uploaded_zip_structure(zf: zipfile.ZipFile) -> list[str]: file_names = [] for member in zf.infolist(): if member.is_dir(): continue member_name = member.filename.replace("\\", "/") if "/" in member_name: raise gr.Error(f"ZIP darf keine Unterordner enthalten: {member_name}") suffix = Path(member_name).suffix.lower() stem = Path(member_name).stem if member_name == "meta.json": file_names.append(member_name) continue if suffix not in ALLOWED_IMAGE_EXTENSIONS: raise gr.Error(f"Nicht erlaubter Dateityp im ZIP: {member_name}") if not (stem.startswith("color_") or stem.startswith("normal_")): raise gr.Error(f"Ungültiger Dateiname im ZIP: {member_name}") file_names.append(member_name) return file_names def overwrite_edit_set_from_zip(zip_file_path: str, session_dir: str): scene_dir = find_single_scene_dir(session_dir) edit_dir = scene_dir / "edit" edit_dir.mkdir(parents=True, exist_ok=True) with zipfile.ZipFile(zip_file_path, "r") as zf: uploaded_names = validate_uploaded_zip_structure(zf) has_uploaded_colors = any(name.startswith("color_") for name in uploaded_names) has_uploaded_normals = any(name.startswith("normal_") for name in uploaded_names) tmp_extract_dir = Path(session_dir) / "_upload_tmp" if tmp_extract_dir.exists(): shutil.rmtree(tmp_extract_dir) tmp_extract_dir.mkdir(parents=True, exist_ok=True) try: zf.extractall(tmp_extract_dir) if has_uploaded_colors and not has_uploaded_normals: for p in edit_dir.glob("normal_*"): if p.is_file(): p.unlink() for name in uploaded_names: if name == "meta.json": continue src = tmp_extract_dir / name dst = edit_dir / name shutil.copy2(src, dst) meta_src = tmp_extract_dir / "meta.json" meta_dst = scene_dir / "meta.json" if meta_src.exists() and not meta_dst.exists(): shutil.copy2(meta_src, meta_dst) finally: if tmp_extract_dir.exists(): shutil.rmtree(tmp_extract_dir) report = inspect_edit_set(edit_dir) if report["state"] == "READY": status = ( "Upload erfolgreich.\n" "Komplettes Fotoset erkannt:\n" f"- Colors: {len(report['found_colors'])}\n" f"- Normalmaps: {len(report['found_normals'])}\n" "Rekonstruktion kann direkt fortgesetzt werden." ) elif report["state"] == "NEEDS_NORMALS": status = ( "Upload erfolgreich.\n" "Nur vollständige Color-Ansichten erkannt, aber keine vollständigen Normalmaps.\n" "Alte Normalmaps wurden verworfen bzw. als ungültig behandelt.\n" "Normalmaps müssen vor der Rekonstruktion neu berechnet werden." ) else: status = ( "Upload unvollständig oder ungültig.\n" f"Fehlende Colors: {report['missing_colors']}\n" f"Fehlende Normalmaps: {report['missing_normals']}\n" "Bitte ein vollständiges Fotoset hochladen." ) preview_paths = sorted(str(p) for p in edit_dir.glob("color_*")) return status, preview_paths def ensure_ready_for_reconstruction(session_dir: str) -> str: scene_dir = find_single_scene_dir(session_dir) edit_dir = scene_dir / "edit" report = inspect_edit_set(edit_dir) if report["state"] == "READY": return "READY" if report["state"] == "NEEDS_NORMALS": raise gr.Error("Es sind nur Color-Bilder vorhanden. Bitte zuerst die Normalmaps neu berechnen.") raise gr.Error( "Das Fotoset ist unvollständig. " f"Fehlende Colors: {report['missing_colors']} | " f"Fehlende Normalmaps: {report['missing_normals']}" ) def run_generate_multiview(session_dir: str): inference_config = "configs/inference-768-6view.yaml" pretrained_model = "./ckpts" crop_size = 740 seed = 600 num_views = 7 save_mode = "rgb" multiview_dir = str(get_multiview_root(session_dir)) subprocess.run( [ sys.executable, "inference.py", "--config", inference_config, f"pretrained_model_name_or_path={pretrained_model}", f"validation_dataset.crop_size={crop_size}", "with_smpl=false", f"validation_dataset.root_dir={session_dir}", f"seed={seed}", f"num_views={num_views}", f"save_mode={save_mode}", "run_mode=generate", f"multiview_tmp_dir={multiview_dir}", "prefer_edited_views=true", ], env=_get_subprocess_env(), check=True ) def run_reconstruct_from_session(session_dir: str): inference_config = "configs/inference-768-6view.yaml" pretrained_model = "./ckpts" crop_size = 740 seed = 600 num_views = 7 save_mode = "rgb" multiview_dir = str(get_multiview_root(session_dir)) subprocess.run( [ sys.executable, "inference.py", "--config", inference_config, f"pretrained_model_name_or_path={pretrained_model}", f"validation_dataset.crop_size={crop_size}", "with_smpl=false", f"validation_dataset.root_dir={session_dir}", f"seed={seed}", f"num_views={num_views}", f"save_mode={save_mode}", "run_mode=reconstruct", f"multiview_tmp_dir={multiview_dir}", "prefer_edited_views=true", ], env=_get_subprocess_env(), check=True ) def collect_outputs_from_session(session_dir: str): scene_dir = find_single_scene_dir(session_dir) scene_name = scene_dir.name output_video = glob(os.path.join("out", scene_name, "*.mp4")) output_objects = glob(os.path.join("out", scene_name, "*.obj")) video = output_video[0] if output_video else None mesh = output_objects[0] if len(output_objects) > 0 else None mesh_color = output_objects[1] if len(output_objects) > 1 else None return video, mesh, mesh_color @spaces.GPU(duration=140) def process_generate(input_pil, remove_bg): torch.cuda.empty_cache() removed_bg_path, session_dir = remove_background(input_pil, remove_bg) try: run_generate_multiview(session_dir) gallery = list_gallery_images(session_dir) status = ( "Stufe 1 abgeschlossen.\n" "Multiview-Bilder wurden erzeugt und im Session-Ordner gespeichert.\n" "Du kannst jetzt das Fotoset herunterladen, extern bearbeiten und wieder hochladen.\n" "Session: " + session_dir ) return session_dir, status, gallery except subprocess.CalledProcessError as e: shutil.rmtree(session_dir, ignore_errors=True) raise gr.Error(f"Fehler während der Multiview-Erzeugung: {str(e)}") finally: torch.cuda.empty_cache() def process_download_set(session_dir): if not session_dir or not Path(session_dir).exists(): raise gr.Error("Kein gültiger Session-Ordner vorhanden.") zip_path = create_edit_zip(session_dir) status = f"Fotoset als ZIP erstellt: {zip_path}" return zip_path, status def process_upload_set(upload_zip, session_dir): zip_path = get_uploaded_file_path(upload_zip) session_dir = ensure_session_from_upload(session_dir, zip_path) status, gallery = overwrite_edit_set_from_zip(zip_path, session_dir) if "Session:" not in status: status = status + f"\nSession: {session_dir}" return session_dir, status, gallery @spaces.GPU(duration=140) def process_reconstruct(session_dir, keep_session): if not session_dir or not Path(session_dir).exists(): raise gr.Error("Kein gültiger Session-Ordner vorhanden.") torch.cuda.empty_cache() try: ensure_ready_for_reconstruction(session_dir) run_reconstruct_from_session(session_dir) video, mesh, mesh_color = collect_outputs_from_session(session_dir) status = "Stufe 2 abgeschlossen. Rekonstruktion erfolgreich." if not keep_session: shutil.rmtree(session_dir, ignore_errors=True) session_dir = "" return status, video, mesh, mesh_color, session_dir except subprocess.CalledProcessError as e: raise gr.Error(f"Fehler während der Rekonstruktion: {str(e)}") finally: torch.cuda.empty_cache() def process_clear_session(session_dir): if session_dir and Path(session_dir).exists(): shutil.rmtree(session_dir, ignore_errors=True) return "", "Session gelöscht.", [], None, None, None, None css = """ div#col-container{ margin: 0 auto; max-width: 1100px; } div#video-out-elm{ height: 323px; } """ def gradio_interface(): with gr.Blocks() as app: with gr.Column(elem_id="col-container"): gr.Markdown("# PSHuman 2.0 – Zwei-Stufen-Pipeline mit Multiview-Export/Import") gr.HTML("""
""") session_dir_box = gr.Textbox(label="Session-Ordner", interactive=False) with gr.Group(): with gr.Row(): with gr.Column(scale=2): input_image = gr.Image( label="Image input", type="pil", image_mode="RGBA", height=480 ) remove_bg = gr.Checkbox(label="Need to remove BG?", value=False) keep_session = gr.Checkbox(label="Session nach Stufe 2 behalten", value=True) btn_generate = gr.Button("1) Multiview erzeugen") btn_download = gr.Button("2) Fotoset herunterladen") upload_zip = gr.File(label="3) Bearbeitetes Fotoset hochladen", file_types=[".zip"]) btn_upload = gr.Button("4) Upload prüfen und Bilder überschreiben") btn_reconstruct = gr.Button("5) Rekonstruktion fortsetzen") btn_clear = gr.Button("Session löschen") with gr.Column(scale=4): status_box = gr.Textbox(label="Status", lines=8) multiview_gallery = gr.Gallery(label="Multiview Edit Set", columns=3, rows=2, height=420) download_file = gr.File(label="Download ZIP") output_video = gr.Video(label="Output Video", elem_id="video-out-elm") with gr.Row(): output_object_mesh = gr.Model3D(label=".OBJ Mesh", height=240) output_object_color = gr.Model3D(label=".OBJ colored", height=240) gr.Examples( examples=examples_folder, inputs=[input_image], examples_per_page=11 ) btn_generate.click( process_generate, inputs=[input_image, remove_bg], outputs=[session_dir_box, status_box, multiview_gallery] ) btn_download.click( process_download_set, inputs=[session_dir_box], outputs=[download_file, status_box] ) btn_upload.click( process_upload_set, inputs=[upload_zip, session_dir_box], outputs=[session_dir_box, status_box, multiview_gallery] ) btn_reconstruct.click( process_reconstruct, inputs=[session_dir_box, keep_session], outputs=[status_box, output_video, output_object_mesh, output_object_color, session_dir_box] ) btn_clear.click( process_clear_session, inputs=[session_dir_box], outputs=[session_dir_box, status_box, multiview_gallery, download_file, output_video, output_object_mesh, output_object_color] ) return app app = gradio_interface() app.launch(show_error=True, ssr_mode=False, css=css)