Spaces:
Sleeping
Sleeping
| """projtex PBR: albedo, roughness and metallic maps from one lit photo of an object. | |
| Qwen-Image-Edit-2511 + the Lightning 4-step LoRA + a LoRA trained (DiffSynth-Studio) on synthetic Blender renders | |
| of MatSynth materials, one instruction per map. API: /decompose (image, steps) -> (albedo, roughness, metallic).""" | |
| import math | |
| import os | |
| import torch | |
| from huggingface_hub import hf_hub_download | |
| from PIL import Image | |
| from safetensors.torch import load_file | |
| try: | |
| import spaces | |
| except ImportError: # outside a ZeroGPU Space (an HF Job testing the same code) | |
| class spaces: # noqa: N801 | |
| def GPU(*a, **k): | |
| return (lambda f: f) if not a or not callable(a[0]) else a[0] | |
| from diffusers import FlowMatchEulerDiscreteScheduler, QwenImageEditPlusPipeline | |
| LORA_REPO = os.environ.get("PBR_LORA_REPO", "Mike0021/qwen-image-edit-2511-pbr-lora") | |
| LORA_FILE = os.environ.get("PBR_LORA_FILE", "pbr2/step-4215.safetensors") | |
| PROMPTS = { | |
| "albedo": "Turn this photo into its albedo map: the base colour of every surface, with all lighting, shadows " | |
| "and reflections removed.", | |
| "roughness": "Turn this photo into its roughness map: grayscale, black where the surface is glossy, white where " | |
| "it is rough.", | |
| "metallic": "Turn this photo into its metallic map: white where the surface is bare metal, black everywhere else.", | |
| } | |
| SCHEDULER = { # Qwen-Image-Lightning's sampling (exponential time shift, shift 3) | |
| "base_image_seq_len": 256, "base_shift": math.log(3), "invert_sigmas": False, "max_image_seq_len": 8192, | |
| "max_shift": math.log(3), "num_train_timesteps": 1000, "shift": 1.0, "shift_terminal": None, | |
| "stochastic_sampling": False, "time_shift_type": "exponential", "use_beta_sigmas": False, | |
| "use_dynamic_shifting": True, "use_exponential_sigmas": False, "use_karras_sigmas": False} | |
| def pbr_lora_state_dict(path: str) -> dict: | |
| """DiffSynth saves PEFT keys of the bare transformer ('...to_q.lora_A.default.weight'); diffusers wants them | |
| under 'transformer.' without the adapter name.""" | |
| sd = load_file(path) | |
| return {"transformer." + k.replace(".default", ""): v for k, v in sd.items()} | |
| pipe = QwenImageEditPlusPipeline.from_pretrained( | |
| "Qwen/Qwen-Image-Edit-2511", scheduler=FlowMatchEulerDiscreteScheduler.from_config(SCHEDULER), | |
| torch_dtype=torch.bfloat16) | |
| pipe.load_lora_weights("lightx2v/Qwen-Image-Edit-2511-Lightning", | |
| weight_name="Qwen-Image-Edit-2511-Lightning-4steps-V1.0-bf16.safetensors", | |
| adapter_name="lightning") | |
| pipe.load_lora_weights(pbr_lora_state_dict(hf_hub_download(LORA_REPO, LORA_FILE, token=os.environ.get("HF_TOKEN"))), | |
| adapter_name="pbr") | |
| names, weights = ["lightning", "pbr"], [1.0, 1.0] | |
| if os.environ.get("PBR_LORA_FILE2"): # a blend of two PBR LoRAs, e.g. PBR_LORA_WEIGHTS="0.5,0.5" | |
| pipe.load_lora_weights(pbr_lora_state_dict(hf_hub_download(LORA_REPO, os.environ["PBR_LORA_FILE2"], | |
| token=os.environ.get("HF_TOKEN"))), adapter_name="pbr2") | |
| w1, w2 = (float(x) for x in os.environ.get("PBR_LORA_WEIGHTS", "0.5,0.5").split(",")) | |
| names, weights = ["lightning", "pbr", "pbr2"], [1.0, w1, w2] | |
| pipe.set_adapters(names, adapter_weights=weights) | |
| pipe.fuse_lora(adapter_names=names, lora_scale=1.0) | |
| pipe.unload_lora_weights() | |
| pipe.to("cuda") | |
| def _square(img: Image.Image, side: int = 768): | |
| w, h = img.size | |
| s = max(w, h) | |
| canvas = Image.new("RGB", (s, s), img.getpixel((0, 0))) | |
| canvas.paste(img, ((s - w) // 2, (s - h) // 2)) | |
| return canvas.resize((side, side), Image.LANCZOS), ((s - w) // 2, (s - h) // 2, w, h, s) | |
| def _unsquare(img: Image.Image, box) -> Image.Image: | |
| x, y, w, h, s = box | |
| return img.resize((s, s), Image.LANCZOS).crop((x, y, x + w, y + h)) | |
| def decompose(image, steps: int = 4): | |
| """Albedo (sRGB), roughness and metallic (grey: 0 black .. 1 white) of the object in ``image``.""" | |
| img = Image.open(image).convert("RGB") if isinstance(image, str) else image.convert("RGB") | |
| sq, box = _square(img) | |
| outs = [] | |
| for key in ("albedo", "roughness", "metallic"): | |
| out = pipe(image=[sq], prompt=PROMPTS[key], negative_prompt=" ", num_inference_steps=int(steps), | |
| true_cfg_scale=1.0, height=768, width=768, | |
| generator=torch.Generator("cuda").manual_seed(0)).images[0] | |
| out = _unsquare(out.convert("RGB"), box) | |
| outs.append(out if key == "albedo" else out.convert("L")) | |
| return tuple(outs) | |
| if __name__ == "__main__" or os.environ.get("SYSTEM") == "spaces": | |
| import gradio as gr | |
| with gr.Blocks(title="projtex PBR") as demo: | |
| gr.Markdown("# PBR maps from a photo of an object\nAlbedo, roughness and metallic in one call: " | |
| "Qwen-Image-Edit-2511 + Lightning + a LoRA trained on renders of MatSynth materials " | |
| "(used by the projection-texturing Space's `pbr` option).") | |
| with gr.Row(): | |
| inp = gr.Image(type="filepath", label="Photo") | |
| steps = gr.Slider(2, 8, value=4, step=1, label="Steps") | |
| btn = gr.Button("Decompose") | |
| with gr.Row(): | |
| a, r, m = gr.Image(label="Albedo"), gr.Image(label="Roughness"), gr.Image(label="Metallic") | |
| btn.click(decompose, [inp, steps], [a, r, m], api_name="decompose") | |
| demo.queue().launch() | |