Spaces:
Runtime error
Runtime error
| """ | |
| MatFuse PBR Material Generator - HuggingFace Space (CPU Version) | |
| """ | |
| import os | |
| import sys | |
| import torch | |
| import numpy as np | |
| from PIL import Image | |
| import gradio as gr | |
| from huggingface_hub import hf_hub_download | |
| from omegaconf import OmegaConf | |
| device = 'cpu' | |
| model = None | |
| def load_model(): | |
| global model | |
| if model is not None: | |
| return model | |
| print("Downloading MatFuse checkpoint...") | |
| ckpt_path = hf_hub_download( | |
| repo_id="gvecchio/MatFuse", | |
| filename="matfuse-pruned.ckpt", | |
| local_dir="checkpoints", | |
| local_dir_use_symlinks=False | |
| ) | |
| print("Cloning MatFuse repository...") | |
| if not os.path.exists('matfuse-sd'): | |
| os.system('git clone https://github.com/giuvecchio/matfuse-sd.git') | |
| sys.path.insert(0, 'matfuse-sd/src') | |
| from ldm.util import instantiate_from_config | |
| config_path = 'matfuse-sd/src/configs/diffusion/matfuse-ldm-vq-4ch.yaml' | |
| config = OmegaConf.load(config_path) | |
| print("Loading model...") | |
| model = instantiate_from_config(config.model) | |
| ckpt = torch.load(ckpt_path, map_location='cpu') | |
| model.load_state_dict(ckpt['state_dict'], strict=False) | |
| model = model.to(device) | |
| model.eval() | |
| print("Model loaded successfully!") | |
| return model | |
| def generate_material(input_image): | |
| global model | |
| try: | |
| if model is None: | |
| model = load_model() | |
| if input_image is None: | |
| blank = Image.new('RGB', (512, 512), (128, 128, 128)) | |
| return blank, blank, blank, blank | |
| input_image = input_image.convert('RGB').resize((512, 512)) | |
| img_np = np.array(input_image).astype(np.float32) / 127.5 - 1.0 | |
| img_tensor = torch.from_numpy(img_np).permute(2, 0, 1).unsqueeze(0).to(device) | |
| if hasattr(model, 'cond_stage_model') and model.cond_stage_model is not None: | |
| cond = model.cond_stage_model.encode(img_tensor) | |
| else: | |
| cond = model.encode_first_stage(img_tensor) | |
| cond = model.get_first_stage_encoding(cond) | |
| samples, _ = model.sample( | |
| cond=cond, | |
| batch_size=1, | |
| return_intermediates=False, | |
| ddim_steps=25, | |
| eta=0.0, | |
| unconditional_guidance_scale=7.5 | |
| ) | |
| outputs = model.decode_first_stage(samples) | |
| outputs = outputs.cpu().numpy() | |
| outputs = ((outputs + 1) * 127.5).clip(0, 255).astype(np.uint8) | |
| if outputs.shape[1] >= 12: | |
| diffuse = Image.fromarray(outputs[0, 0:3].transpose(1, 2, 0)) | |
| normal = Image.fromarray(outputs[0, 3:6].transpose(1, 2, 0)) | |
| roughness = Image.fromarray(outputs[0, 6:9].transpose(1, 2, 0)) | |
| specular = Image.fromarray(outputs[0, 9:12].transpose(1, 2, 0)) | |
| elif outputs.shape[1] == 3: | |
| normal = Image.fromarray(outputs[0].transpose(1, 2, 0)) | |
| diffuse = input_image | |
| roughness = Image.new('RGB', (512, 512), (128, 128, 128)) | |
| specular = Image.new('RGB', (512, 512), (64, 64, 64)) | |
| else: | |
| diffuse = input_image | |
| normal = Image.new('RGB', (512, 512), (128, 128, 255)) | |
| roughness = Image.new('RGB', (512, 512), (128, 128, 128)) | |
| specular = Image.new('RGB', (512, 512), (64, 64, 64)) | |
| return diffuse, normal, roughness, specular | |
| except Exception as e: | |
| import traceback | |
| traceback.print_exc() | |
| blank = Image.new('RGB', (512, 512), (0, 0, 0)) | |
| return blank, blank, blank, blank | |
| demo = gr.Interface( | |
| fn=generate_material, | |
| inputs=gr.Image(type="pil", label="Input Image"), | |
| outputs=[ | |
| gr.Image(type="pil", label="Diffuse/Albedo"), | |
| gr.Image(type="pil", label="Normal"), | |
| gr.Image(type="pil", label="Roughness"), | |
| gr.Image(type="pil", label="Specular") | |
| ], | |
| title="MatFuse PBR Material Generator", | |
| description="Upload an image to generate PBR texture maps. Running on CPU - first request takes 5-10 min (model loading), then 2-5 min per image.", | |
| allow_flagging="never" | |
| ) | |
| demo.launch() | |