SculptOk-Mobile / app.py
TS447's picture
Update app.py
1289d31 verified
import gradio as gr
import torch
import numpy as np
from diffusers import MarigoldDepthPipeline
from PIL import Image
import tempfile
import os
print("Loading Marigold Diffusion Model...")
# Marigold LCM load ho raha hai (CPU par)
pipe = MarigoldDepthPipeline.from_pretrained(
"prs-eth/marigold-depth-lcm-v1-0",
torch_dtype=torch.float32
)
def generate_pro_depth(image, invert_depth):
# AI Engine Processing
output = pipe(
image,
num_inference_steps=4,
ensemble_size=1
)
depth_map = output.prediction
depth_map = np.squeeze(depth_map)
# Normalization (0 se 1 ke beech value lana)
depth_min = depth_map.min()
depth_max = depth_map.max()
depth_normalized = (depth_map - depth_min) / (depth_max - depth_min)
# ---------------------------------------------------------
# 1. PREVIEW KE LIYE (8-Bit) - Taki screen par safed dabba na aaye
# ---------------------------------------------------------
depth_8bit = (depth_normalized * 255).astype(np.uint8)
if invert_depth:
depth_8bit = 255 - depth_8bit
preview_image = Image.fromarray(depth_8bit, mode='L')
# ---------------------------------------------------------
# 2. CNC ARTCAM KE LIYE (16-Bit) - Asli smooth file
# ---------------------------------------------------------
# 16-bit mein 0 se 65535 tak values hoti hain
depth_16bit = (depth_normalized * 65535).astype(np.uint16)
if invert_depth:
depth_16bit = 65535 - depth_16bit
# Isko 'I;16' mode mein save karna zaroori hai 16-bit PNG ke liye
img_16bit = Image.fromarray(depth_16bit, mode='I;16')
# Ek temporary file banakar save kar rahe hain taki tu download kar sake
temp_dir = tempfile.gettempdir()
file_path = os.path.join(temp_dir, "Taha_Industry_16bit_Depth.png")
img_16bit.save(file_path)
# Dono chizein return kar rahe hain: Preview Image aur Downloadable File
return preview_image, file_path
# Gradio Interface Setup
iface = gr.Interface(
fn=generate_pro_depth,
inputs=[
gr.Image(type="pil", label="Upload 2D Image"),
gr.Checkbox(label="Invert Depth (Emboss Effect)", value=True)
],
outputs=[
gr.Image(type="pil", label="Preview (8-bit)"),
gr.File(label="Download 16-bit HD File for ARTCAM") # Ye naya download button hai
],
title="Taha Industry Pro CNC API (16-Bit Enabled)",
description="Shows an 8-bit preview on screen but provides a raw 16-bit PNG file download for high-precision CNC carving."
)
iface.launch()