""" MatFuse — PBR Material Generation Demo Gradio app for generating physically-based rendering (PBR) material maps using the MatFuse diffusion model. Supports text, image, sketch, and color-palette conditioning. Designed for Hugging Face Spaces with ZeroGPU support. """ import os import random from typing import Optional import gradio as gr import numpy as np import spaces import torch from diffusers import DiffusionPipeline from PIL import Image # --------------------------------------------------------------------------- # Model loading # --------------------------------------------------------------------------- REPO_ID = os.environ.get("MATFUSE_REPO", "gvecchio/MatFuse") pipe = DiffusionPipeline.from_pretrained( REPO_ID, trust_remote_code=True, torch_dtype=torch.float16, ) # --------------------------------------------------------------------------- # Palette extraction (lightweight K-Means, no heavy deps) # --------------------------------------------------------------------------- def extract_palette(image: Image.Image, n_colors: int = 5) -> list[list[int]]: """Extract dominant colors from an image using simple K-Means.""" img = image.convert("RGB").resize((64, 64)) pixels = np.array(img).reshape(-1, 3).astype(np.float32) # Mini K-Means rng = np.random.default_rng(0) centroids = pixels[rng.choice(len(pixels), n_colors, replace=False)] for _ in range(20): dists = np.linalg.norm(pixels[:, None] - centroids[None], axis=2) labels = dists.argmin(axis=1) for k in range(n_colors): mask = labels == k if mask.any(): centroids[k] = pixels[mask].mean(axis=0) return centroids.clip(0, 255).astype(np.uint8).tolist() def palette_to_image(colors: list[list[int]], height: int = 50) -> Image.Image: """Render a palette swatch strip.""" n = len(colors) w_each = 60 img = Image.new("RGB", (w_each * n, height)) for i, c in enumerate(colors): for x in range(w_each): for y in range(height): img.putpixel((i * w_each + x, y), tuple(c)) return img # --------------------------------------------------------------------------- # Generation # --------------------------------------------------------------------------- @spaces.GPU @torch.inference_mode() def generate( prompt: Optional[str], image: Optional[Image.Image], palette_image: Optional[Image.Image], sketch: Optional[Image.Image], guidance_scale: float, num_steps: int, seed: int, randomize_seed: bool, ): """Run the MatFuse pipeline and return the four PBR maps + palette preview.""" if randomize_seed: seed = random.randint(0, 2**31 - 1) # Move to GPU (ZeroGPU allocates on call) pipe.to("cuda") # --- Build kwargs ------------------------------------------------------- kwargs: dict = dict( num_inference_steps=num_steps, guidance_scale=guidance_scale, generator=torch.Generator("cuda").manual_seed(seed), ) # Text if prompt and prompt.strip(): kwargs["text"] = prompt.strip() # Reference image if image is not None: kwargs["image"] = image # Sketch if sketch is not None: kwargs["sketch"] = sketch # Palette (extracted from an uploaded image) palette_preview = None if palette_image is not None: colors = extract_palette(palette_image, n_colors=5) palette_arr = np.array(colors, dtype=np.float32) / 255.0 kwargs["palette"] = palette_arr palette_preview = palette_to_image(colors) result = pipe(**kwargs) diffuse_img = result["diffuse"][0] normal_img = result["normal"][0] roughness_img = result["roughness"][0] specular_img = result["specular"][0] return diffuse_img, normal_img, roughness_img, specular_img, palette_preview, seed # --------------------------------------------------------------------------- # Example data # --------------------------------------------------------------------------- EXAMPLE_PROMPTS = [ "Red brick wall with white mortar", "Polished oak wood floor", "Rough concrete with cracks", "Mossy cobblestone path", "Shiny marble tiles", "Rusted metal panel", ] # --------------------------------------------------------------------------- # UI # --------------------------------------------------------------------------- css = """ #matfuse-title { text-align: center; margin-bottom: 0.5em; } #matfuse-subtitle { text-align: center; color: #666; margin-top: 0; } .output-map img { border-radius: 8px; } footer { display: none !important; } """ with gr.Blocks(title="MatFuse — PBR Material Generator") as demo: # Header gr.Markdown("# MatFuse", elem_id="matfuse-title") gr.Markdown( "Generate seamless PBR material maps (diffuse, normal, roughness, specular) " "from text, images, sketches, and color palettes. " "[Paper](https://arxiv.org/abs/2308.11408)  |  " "[Code](https://github.com/gvecchio/matfuse-sd)", elem_id="matfuse-subtitle", ) with gr.Row(): # ── Left column: inputs ────────────────────────────────────── with gr.Column(scale=1): prompt = gr.Textbox( label="Text prompt", placeholder="e.g. 'Old wooden floor with scratches'", lines=2, ) with gr.Accordion("Image conditioning", open=False): image_input = gr.Image( label="Reference image", type="pil", sources=["upload", "clipboard"], ) with gr.Accordion("Palette conditioning", open=False): palette_image = gr.Image( label="Upload image to extract palette", type="pil", sources=["upload", "clipboard"], ) with gr.Accordion("Sketch conditioning", open=False): sketch_input = gr.Image( label="Binary sketch / edge map", type="pil", image_mode="L", sources=["upload", "clipboard"], ) with gr.Accordion("Generation settings", open=False): guidance_scale = gr.Slider( label="Guidance scale", minimum=1.0, maximum=15.0, value=4.0, step=0.5, ) num_steps = gr.Slider( label="Inference steps", minimum=10, maximum=100, value=50, step=5, ) with gr.Row(): seed = gr.Number(label="Seed", value=42, precision=0) randomize_seed = gr.Checkbox(label="Randomize", value=True) generate_btn = gr.Button("Generate", variant="primary", size="lg") gr.Examples( examples=[[p] for p in EXAMPLE_PROMPTS], inputs=[prompt], label="Example prompts", ) # ── Right column: outputs ──────────────────────────────────── with gr.Column(scale=1): with gr.Row(): diffuse_out = gr.Image(label="Diffuse", elem_classes="output-map", interactive=False) normal_out = gr.Image(label="Normal", elem_classes="output-map", interactive=False) with gr.Row(): roughness_out = gr.Image(label="Roughness", elem_classes="output-map", interactive=False) specular_out = gr.Image(label="Specular", elem_classes="output-map", interactive=False) palette_out = gr.Image(label="Extracted palette", visible=True, height=60, interactive=False) seed_out = gr.Number(label="Seed used", interactive=False) # ── Wiring ────────────────────────────────────────────────────── generate_btn.click( fn=generate, inputs=[ prompt, image_input, palette_image, sketch_input, guidance_scale, num_steps, seed, randomize_seed, ], outputs=[diffuse_out, normal_out, roughness_out, specular_out, palette_out, seed_out], ) if __name__ == "__main__": demo.launch(css=css, theme=gr.themes.Soft()) # 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