Spaces:
Running on Zero
Running on Zero
| """FlashReFocus for the Hugging Face Space (ZeroGPU), on the flashrefocus package (a copy of the repository's). | |
| ZeroGPU gives a GPU only inside @spaces.GPU calls: every model is built on the CPU, moved to `cuda` at import time as an | |
| ordinary module (Stage I with two separate U-Nets, no weights swapped in place), and the per-session image and depth | |
| stay on the CPU between calls. The weights are downloaded on start; the FlashReFocus ones come from a private model | |
| repository read with the HF_TOKEN secret.""" | |
| import os | |
| from pathlib import Path | |
| import numpy as np | |
| import torch | |
| from huggingface_hub import hf_hub_download, snapshot_download | |
| from PIL import Image | |
| HERE = Path(__file__).resolve().parent | |
| WEIGHT_REPO = os.environ.get("FLASHREFOCUS_WEIGHT_REPO", "ming0531/FlashReFocus-weights") | |
| def fetch_weights(): | |
| """Weights into ./weights (skipped when FLASHREFOCUS_WEIGHTS already points to a local copy).""" | |
| if "FLASHREFOCUS_WEIGHTS" in os.environ: | |
| return | |
| w = HERE / "weights" | |
| sd_files = ["model_index.json", "scheduler/scheduler_config.json", "tokenizer/*", "text_encoder/config.json", | |
| "text_encoder/model.safetensors", "unet/config.json", "unet/diffusion_pytorch_model.safetensors", | |
| "vae/config.json", "vae/diffusion_pytorch_model.safetensors"] | |
| snapshot_download("sd2-community/stable-diffusion-2-1-base", allow_patterns=sd_files, | |
| local_dir=w / "stable-diffusion-2-1-base") | |
| snapshot_download("stable-diffusion-v1-5/stable-diffusion-v1-5", | |
| allow_patterns=["vae/config.json", "vae/diffusion_pytorch_model.safetensors"], | |
| local_dir=w / "stable-diffusion-v1-5") | |
| hf_hub_download("depth-anything/Depth-Anything-V2-Large", "depth_anything_v2_vitl.pth", local_dir=w) | |
| for f in ("stage1.pkl", "refiner.pt", "renderer.pt"): | |
| hf_hub_download(WEIGHT_REPO, f, local_dir=w, token=os.environ.get("HF_TOKEN")) | |
| os.environ["FLASHREFOCUS_WEIGHTS"] = str(w) | |
| fetch_weights() | |
| os.environ.setdefault("FLASHREFOCUS_BUILD_DEVICE", "cpu") # third_party/PiSA-SR here builds on this device | |
| from flashrefocus import Refocuser, Scene, paths # noqa: E402 | |
| from flashrefocus.deblurring import Deblurrer # noqa: E402 | |
| from flashrefocus.image_io import image_to_tensor, tensor_to_uint8 # noqa: E402 | |
| from flashrefocus.rendering import DepthEstimator, depth_edits, load_renderer # noqa: E402 | |
| def build_refocuser(): | |
| deblurrer = Deblurrer.from_pretrained(separate_unets=True, device="cpu") | |
| del deblurrer.model.text_encoder # the empty-prompt embedding is computed once, at construction | |
| return Refocuser(deblurrer.to("cuda"), DepthEstimator.from_pretrained(device="cpu").to("cuda"), | |
| load_renderer(paths.RENDERER, "cpu").to("cuda")) | |
| refocuser = build_refocuser() | |
| def restore(rgb, omega): | |
| """HxWx3 uint8 (sides multiples of 16) -> restored image (HxWx3 uint8) and its depth (HxW float32, 1 = far).""" | |
| refocuser.deblurrer.detail_weight = float(omega) | |
| scene = refocuser.prepare(image_to_tensor(rgb).cuda()) | |
| return tensor_to_uint8(scene.image), scene.depth[0].cpu().numpy() | |
| def render(rgb, depth, focal_depth, strength): | |
| """Bokeh (HxWx3 uint8) of the restored image with blur strength K and focal depth D_f; depth may be edited.""" | |
| d = torch.from_numpy(np.ascontiguousarray(depth, dtype=np.float32))[None].cuda() | |
| return tensor_to_uint8(refocuser.render(Scene(image_to_tensor(rgb).cuda(), d), strength, focal_depth)) | |
| def two_focal_planes(depth, focal_depth_a, focal_depth_b): | |
| return depth_edits.two_focal_planes(torch.from_numpy(depth)[None], focal_depth_a, focal_depth_b)[0].numpy() | |
| def tilt_shift(height, width, row, band=0.06): | |
| return depth_edits.tilt_shift(height, width, row, band=band)[0].numpy() | |
| def depth_from_image(image, height, width): | |
| """A user depth map (near is bright, any size) -> D in [0, 1], 1 = far.""" | |
| grey = Image.fromarray(np.asarray(image)).convert("L").resize((width, height), Image.BILINEAR) | |
| return (1.0 - np.asarray(grey, dtype=np.float32) / 255.0).astype(np.float32) | |