Spaces:
Running on Zero
Running on Zero
| import gradio as gr | |
| import numpy as np | |
| from PIL import Image | |
| import cv2 | |
| import torch | |
| import zipfile | |
| import io | |
| import tempfile | |
| import spaces | |
| from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler | |
| print(f"Gradio version: {gr.__version__}", flush=True) | |
| # ─── Model ──────────────────────────────────────────────────────────────────── | |
| pipe = None | |
| def load_model(): | |
| global pipe | |
| if pipe is None: | |
| print("Loading SD 1.5...", flush=True) | |
| pipe = StableDiffusionPipeline.from_pretrained( | |
| "runwayml/stable-diffusion-v1-5", | |
| torch_dtype=torch.float16, | |
| safety_checker=None, | |
| requires_safety_checker=False, | |
| ) | |
| pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config) | |
| pipe = pipe.to("cuda") | |
| pipe.enable_attention_slicing() | |
| print("Model ready.", flush=True) | |
| return pipe | |
| # ─── PBR Map Generation (CPU/OpenCV) ───────────────────────────────────────── | |
| def make_seamless(img): | |
| arr = np.array(img); h, w = arr.shape[:2]; r = arr.copy(); b = min(h,w)//4 | |
| for x in range(b): a=x/b; r[:,x]=(arr[:,x]*a+arr[:,w-b+x]*(1-a)).astype(np.uint8) | |
| for y in range(b): a=y/b; r[y,:]=(r[y,:]*a+r[h-b+y,:]*(1-a)).astype(np.uint8) | |
| return Image.fromarray(r) | |
| def gen_normal(img): | |
| g = cv2.GaussianBlur(np.array(img.convert("L")).astype(np.float32)/255,(3,3),0) | |
| dx = cv2.Sobel(g,cv2.CV_32F,1,0,ksize=3)*4 | |
| dy = cv2.Sobel(g,cv2.CV_32F,0,1,ksize=3)*4 | |
| dz = np.ones_like(dx); L = np.sqrt(dx**2+dy**2+dz**2) | |
| return Image.fromarray(np.stack([ | |
| ((-dx/L+1)/2*255).astype(np.uint8), | |
| ((-dy/L+1)/2*255).astype(np.uint8), | |
| ((dz/L)*255).astype(np.uint8) | |
| ], axis=-1), mode="RGB") | |
| def gen_roughness(img): | |
| g = np.array(img.convert("L")).astype(np.float32)/255 | |
| return Image.fromarray(((np.abs(g-cv2.GaussianBlur(g,(5,5),0))*8).clip(.2,1)*255).astype(np.uint8), mode="L") | |
| def gen_metallic(img): | |
| hsv = cv2.cvtColor(np.array(img).astype(np.float32)/255, cv2.COLOR_RGB2HSV) | |
| return Image.fromarray(((hsv[:,:,2]*(1-hsv[:,:,1])*.6).clip(0,1)*255).astype(np.uint8), mode="L") | |
| def gen_ao(img): | |
| g = np.array(img.convert("L")).astype(np.float32)/255 | |
| ao = (g/(cv2.GaussianBlur(g,(31,31),0)+.01)).clip(0,1) | |
| ao = cv2.normalize(ao,None,0,255,cv2.NORM_MINMAX).astype(np.uint8) | |
| return Image.fromarray((ao.astype(np.float32)*.85+20).clip(0,255).astype(np.uint8), mode="L") | |
| def gen_height(img): | |
| g = cv2.GaussianBlur(np.array(img.convert("L")),(3,3),0) | |
| return Image.fromarray(cv2.createCLAHE(clipLimit=2,tileGridSize=(8,8)).apply(g), mode="L") | |
| def all_maps(albedo, seamless): | |
| if seamless: albedo = make_seamless(albedo) | |
| return albedo, gen_normal(albedo), gen_roughness(albedo), gen_metallic(albedo), gen_ao(albedo), gen_height(albedo) | |
| def to_zip(maps): | |
| tmp = tempfile.NamedTemporaryFile(suffix=".zip", delete=False) | |
| with zipfile.ZipFile(tmp.name,"w",zipfile.ZIP_DEFLATED) as zf: | |
| for name, img in maps.items(): | |
| buf = io.BytesIO(); img.save(buf, format="PNG") | |
| zf.writestr(f"{name}.png", buf.getvalue()) | |
| return tmp.name | |
| # ─── Gradio Functions ───────────────────────────────────────────────────────── | |
| STYLES = { | |
| "None": "", | |
| "Stone / Rock": "stone texture, rocky surface, seamless", | |
| "Wood": "wood grain texture, natural planks, seamless", | |
| "Metal": "brushed metal surface, industrial, seamless", | |
| "Fabric / Cloth": "woven fabric texture, seamless", | |
| "Concrete": "concrete cement surface, seamless", | |
| "Brick": "brick wall, masonry, seamless", | |
| "Ground / Dirt": "dirt soil ground, natural, seamless", | |
| "Sci-Fi / Tech": "sci-fi futuristic panel, seamless", | |
| } | |
| def from_prompt(prompt, neg, style, res, steps, guidance, seamless, seed): | |
| full = f"{prompt}, {STYLES.get(style,'')}, seamless texture, top-down, PBR, flat lighting, photorealistic".strip(", ") | |
| neg_full = f"{neg}, shadows, 3d render, person, face, vignette" | |
| gen = torch.Generator(device="cuda").manual_seed(int(seed)) if seed >= 0 else None | |
| albedo = load_model()( | |
| prompt=full, negative_prompt=neg_full, | |
| width=int(res), height=int(res), | |
| num_inference_steps=int(steps), | |
| guidance_scale=float(guidance), | |
| generator=gen, | |
| ).images[0] | |
| a,n,r,m,ao,h = all_maps(albedo, seamless) | |
| return a, n, r, m, ao, h, to_zip({"T_Albedo":a,"T_Normal":n,"T_Roughness":r,"T_Metallic":m,"T_AO":ao,"T_Height":h}) | |
| def from_image(img, seamless): | |
| if img is None: | |
| raise gr.Error("Please upload an image.") | |
| a,n,r,m,ao,h = all_maps(img, seamless) | |
| return a, n, r, m, ao, h, to_zip({"T_Albedo":a,"T_Normal":n,"T_Roughness":r,"T_Metallic":m,"T_AO":ao,"T_Height":h}) | |
| # ─── UI (Gradio 4.x compatible) ─────────────────────────────────────────────── | |
| with gr.Blocks(title="TextureForge — PBR Texture Generator") as demo: | |
| gr.Markdown(""" | |
| # 🎨 TextureForge — PBR Texture Generator | |
| Generate full **PBR map sets** (Albedo · Normal · Roughness · Metallic · AO · Height) ready for **Unity3D**. | |
| """) | |
| with gr.Tabs(): | |
| # ── Tab 1: Text → Texture ────────────────────────────────────────────── | |
| with gr.Tab("Text to Texture"): | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| t_prompt = gr.Textbox( | |
| label="Texture Description", | |
| placeholder="weathered oak floorboards with visible grain and knots...", | |
| lines=3 | |
| ) | |
| t_neg = gr.Textbox( | |
| label="Negative Prompt", | |
| value="blurry, low quality, watermark, text, 3d render", | |
| lines=2 | |
| ) | |
| t_style = gr.Dropdown( | |
| label="Style Preset", | |
| choices=list(STYLES.keys()), | |
| value="None" | |
| ) | |
| with gr.Row(): | |
| t_res = gr.Dropdown(label="Resolution", choices=[256,512,768,1024], value=512) | |
| t_steps = gr.Slider(label="Steps", minimum=15, maximum=50, step=1, value=25) | |
| with gr.Row(): | |
| t_guid = gr.Slider(label="Guidance Scale", minimum=3.0, maximum=15.0, step=0.5, value=7.5) | |
| t_seed = gr.Number(label="Seed (-1 = random)", value=-1) | |
| t_seam = gr.Checkbox(label="Make Seamless / Tileable", value=True) | |
| t_btn = gr.Button("🎨 Generate PBR Maps", variant="primary") | |
| with gr.Column(scale=2): | |
| with gr.Row(): | |
| o_alb = gr.Image(label="Albedo — T_Albedo", type="pil") | |
| o_nrm = gr.Image(label="Normal — T_Normal", type="pil") | |
| o_rgh = gr.Image(label="Roughness — T_Roughness", type="pil") | |
| with gr.Row(): | |
| o_met = gr.Image(label="Metallic — T_Metallic", type="pil") | |
| o_ao = gr.Image(label="AO — T_AO", type="pil") | |
| o_hgt = gr.Image(label="Height — T_Height", type="pil") | |
| o_zip = gr.File(label="⬇ Download All Maps (ZIP — Unity-ready)") | |
| t_btn.click( | |
| fn=from_prompt, | |
| inputs=[t_prompt, t_neg, t_style, t_res, t_steps, t_guid, t_seam, t_seed], | |
| outputs=[o_alb, o_nrm, o_rgh, o_met, o_ao, o_hgt, o_zip] | |
| ) | |
| # ── Tab 2: Image → PBR ───────────────────────────────────────────────── | |
| with gr.Tab("Image to PBR Maps"): | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| i_img = gr.Image(label="Upload Base Texture / Albedo", type="pil") | |
| i_seam = gr.Checkbox(label="Make Seamless / Tileable", value=True) | |
| i_btn = gr.Button("🔬 Extract PBR Maps", variant="primary") | |
| gr.Markdown(""" | |
| **TIP:** Upload any real-world surface photo. | |
| Works best with flat, evenly-lit, top-down shots. | |
| **Unity map names:** | |
| `T_Albedo` · `T_Normal` · `T_Roughness` | |
| `T_Metallic` · `T_AO` · `T_Height` | |
| """) | |
| with gr.Column(scale=2): | |
| with gr.Row(): | |
| i_alb = gr.Image(label="Albedo — T_Albedo", type="pil") | |
| i_nrm = gr.Image(label="Normal — T_Normal", type="pil") | |
| i_rgh = gr.Image(label="Roughness — T_Roughness", type="pil") | |
| with gr.Row(): | |
| i_met = gr.Image(label="Metallic — T_Metallic", type="pil") | |
| i_ao = gr.Image(label="AO — T_AO", type="pil") | |
| i_hgt = gr.Image(label="Height — T_Height", type="pil") | |
| i_zip = gr.File(label="⬇ Download All Maps (ZIP — Unity-ready)") | |
| i_btn.click( | |
| fn=from_image, | |
| inputs=[i_img, i_seam], | |
| outputs=[i_alb, i_nrm, i_rgh, i_met, i_ao, i_hgt, i_zip] | |
| ) | |
| # ── Tab 3: Unity Guide ───────────────────────────────────────────────── | |
| with gr.Tab("Unity3D Setup Guide"): | |
| gr.Markdown(""" | |
| ## Using TextureForge Maps in Unity3D | |
| ### 1. Extract the ZIP | |
| Unzip into `Assets/Textures/<MaterialName>/` | |
| ### 2. Texture Import Settings | |
| | File | Texture Type | sRGB | | |
| |------|-------------|------| | |
| | T_Albedo | Default | ✅ ON | | |
| | T_Normal | **Normal Map** | ❌ OFF | | |
| | T_Roughness | Default | ❌ OFF | | |
| | T_Metallic | Default | ❌ OFF | | |
| | T_AO | Default | ❌ OFF | | |
| | T_Height | Default | ❌ OFF | | |
| ### 3. Create URP/Lit Material | |
| **Shader:** `Universal Render Pipeline/Lit` | |
| - Base Map → `T_Albedo` | |
| - Normal Map → `T_Normal` | |
| - Metallic Map → `T_Metallic` | |
| - Occlusion → `T_AO` | |
| ### 4. Roughness → Smoothness | |
| Unity uses **Smoothness** (= 1 − Roughness). | |
| Use the included **TextureForgeHDRPPacker** Unity Editor script to pack channels automatically: | |
| `Tools → TextureForge → Pack HDRP Mask Map` | |
| ### 5. Tiling | |
| Set material Tiling (e.g. `X: 4, Y: 4`) for large meshes since all textures are seamless. | |
| > **HDRP Mask Map:** Pack R=Metallic, G=AO, B=Detail, A=Smoothness using the Unity Editor tool included in the download. | |
| """) | |
| demo.launch() | |