PolyGenixAI6.0 / app.py
anvilinteractiv's picture
Update app.py
15f252d verified
raw
history blame
25.9 kB
import spaces
import os
import gradio as gr
import numpy as np
import torch
from torch.cuda.amp import autocast
import trimesh
import random
from PIL import Image
from transformers import AutoModelForImageSegmentation
from torchvision import transforms
from huggingface_hub import hf_hub_download, snapshot_download
import subprocess
import shutil
import base64
import logging
import time
import traceback
import requests
# Set up logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)
# Install additional dependencies
try:
subprocess.run("pip install spandrel==0.4.1 --no-deps", shell=True, check=True)
except Exception as e:
logger.error(f"Failed to install spandrel: {str(e)}\n{traceback.format_exc()}")
raise
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DTYPE = torch.float16
logger.info(f"Using device: {DEVICE}")
DEFAULT_FACE_NUMBER = 20000 # Reduced for memory efficiency
MAX_SEED = np.iinfo(np.int32).max
TRIPOSG_REPO_URL = "https://github.com/VAST-AI-Research/TripoSG.git"
MV_ADAPTER_REPO_URL = "https://github.com/huanngzh/MV-Adapter.git"
RMBG_PRETRAINED_MODEL = "checkpoints/RMBG-1.4"
TRIPOSG_PRETRAINED_MODEL = "checkpoints/TripoSG"
TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "tmp")
os.makedirs(TMP_DIR, exist_ok=True)
TRIPOSG_CODE_DIR = "./triposg"
if not os.path.exists(TRIPOSG_CODE_DIR):
logger.info(f"Cloning TripoSG repository to {TRIPOSG_CODE_DIR}")
os.system(f"git clone {TRIPOSG_REPO_URL} {TRIPOSG_CODE_DIR}")
MV_ADAPTER_CODE_DIR = "./mv_adapter"
if not os.path.exists(MV_ADAPTER_CODE_DIR):
logger.info(f"Cloning MV-Adapter repository to {MV_ADAPTER_CODE_DIR}")
os.system(f"git clone {MV_ADAPTER_REPO_URL} {MV_ADAPTER_CODE_DIR} && cd {MV_ADAPTER_CODE_DIR} && git checkout 7d37a97e9bc223cdb8fd26a76bd8dd46504c7c3d")
import sys
sys.path.append(TRIPOSG_CODE_DIR)
sys.path.append(os.path.join(TRIPOSG_CODE_DIR, "scripts"))
sys.path.append(MV_ADAPTER_CODE_DIR)
sys.path.append(os.path.join(MV_ADAPTER_CODE_DIR, "scripts"))
try:
from image_process import prepare_image
from briarmbg import BriaRMBG
snapshot_download("briaai/RMBG-1.4", local_dir=RMBG_PRETRAINED_MODEL)
rmbg_net = BriaRMBG.from_pretrained(RMBG_PRETRAINED_MODEL).to(DEVICE, dtype=DTYPE)
rmbg_net.eval()
from triposg.pipelines.pipeline_triposg import TripoSGPipeline
snapshot_download("VAST-AI/TripoSG", local_dir=TRIPOSG_PRETRAINED_MODEL)
triposg_pipe = TripoSGPipeline.from_pretrained(TRIPOSG_PRETRAINED_MODEL).to(DEVICE, dtype=DTYPE)
except Exception as e:
logger.error(f"Failed to load TripoSG models: {str(e)}\n{traceback.format_exc()}")
raise
try:
NUM_VIEWS = 4 # Reduced for memory efficiency
from inference_ig2mv_sdxl import prepare_pipeline, preprocess_image, remove_bg
from mvadapter.utils import get_orthogonal_camera, tensor_to_image, make_image_grid
from mvadapter.utils.render import NVDiffRastContextWrapper, load_mesh, render
mv_adapter_pipe = prepare_pipeline(
base_model="stabilityai/stable-diffusion-xl-base-1.0",
vae_model="madebyollin/sdxl-vae-fp16-fix",
unet_model=None,
lora_model=None,
adapter_path="huanngzh/mv-adapter",
scheduler=None,
num_views=NUM_VIEWS,
device=DEVICE,
dtype=torch.float16,
)
birefnet = AutoModelForImageSegmentation.from_pretrained(
"ZhengPeng7/BiRefNet", trust_remote_code=True
).to(DEVICE, dtype=DTYPE)
transform_image = transforms.Compose(
[
transforms.Resize((512, 512)), # Reduced resolution
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
]
)
remove_bg_fn = lambda x: remove_bg(x, birefnet, transform_image, DEVICE)
except Exception as e:
logger.error(f"Failed to load MV-Adapter models: {str(e)}\n{traceback.format_exc()}")
raise
try:
if not os.path.exists("checkpoints/RealESRGAN_x2plus.pth"):
hf_hub_download("dtarnow/UPscaler", filename="RealESRGAN_x2plus.pth", local_dir="checkpoints")
if not os.path.exists("checkpoints/big-lama.pt"):
subprocess.run("wget -P checkpoints/ https://github.com/Sanster/models/releases/download/add_big_lama/big-lama.pt", shell=True, check=True)
except Exception as e:
logger.error(f"Failed to download checkpoints: {str(e)}\n{traceback.format_exc()}")
raise
def log_gpu_memory():
if torch.cuda.is_available():
allocated = torch.cuda.memory_allocated() / 1024**3
reserved = torch.cuda.memory_reserved() / 1024**3
logger.info(f"GPU Memory: Allocated {allocated:.2f} GB, Reserved {reserved:.2f} GB")
def get_random_hex():
random_bytes = os.urandom(8)
random_hex = random_bytes.hex()
return random_hex
def retry_on_failure(func, max_attempts=3, delay=1):
for attempt in range(max_attempts):
try:
return func()
except RuntimeError as e:
logger.warning(f"Attempt {attempt + 1} failed: {str(e)}\n{traceback.format_exc()}")
if attempt == max_attempts - 1:
raise
time.sleep(delay)
@spaces.GPU(duration=2)
@torch.no_grad()
def run_segmentation(image):
try:
log_gpu_memory()
if isinstance(image, dict):
image_path = image.get("path") or image.get("url")
if not image_path:
raise ValueError("Invalid image input: no path or URL provided")
if image_path.startswith("http"):
temp_image_path = os.path.join(TMP_DIR, f"input_{get_random_hex()}.png")
image_path = download_image(image_path, temp_image_path)
elif isinstance(image, str) and image.startswith("http"):
temp_image_path = os.path.join(TMP_DIR, f"input_{get_random_hex()}.png")
image_path = download_image(image, temp_image_path)
else:
image_path = image
if not isinstance(image, (str, bytes)) or (isinstance(image, str) and not os.path.exists(image)):
raise ValueError(f"Expected str (path/URL), bytes, or FileData dict, got {type(image)}")
with autocast():
image_seg = prepare_image(image_path, bg_color=np.array([1.0, 1.0, 1.0]), rmbg_net=rmbg_net)
rmbg_net.to("cpu")
torch.cuda.empty_cache()
log_gpu_memory()
return image_seg
except Exception as e:
logger.error(f"Error in run_segmentation: {str(e)}\n{traceback.format_exc()}")
raise
@spaces.GPU(duration=3)
@torch.no_grad()
def image_to_3d(image, seed, num_inference_steps=30, guidance_scale=7.0, simplify=True, target_face_num=DEFAULT_FACE_NUMBER, req=None):
try:
log_gpu_memory()
triposg_pipe.to(DEVICE, dtype=DTYPE)
with autocast():
outputs = triposg_pipe(
image=image,
generator=torch.Generator(device=triposg_pipe.device).manual_seed(seed),
num_inference_steps=num_inference_steps,
guidance_scale=guidance_scale
).samples[0]
mesh = trimesh.Trimesh(outputs[0].astype(np.float32), np.ascontiguousarray(outputs[1]))
if simplify:
from utils import simplify_mesh
mesh = simplify_mesh(mesh, target_face_num)
save_dir = os.path.join(TMP_DIR, str(req.session_hash) if req else "examples")
os.makedirs(save_dir, exist_ok=True)
mesh_path = os.path.join(save_dir, f"polygenixai_{get_random_hex()}.glb")
mesh.export(mesh_path)
triposg_pipe.to("cpu")
torch.cuda.empty_cache()
log_gpu_memory()
return mesh_path
except Exception as e:
logger.error(f"Error in image_to_3d: {str(e)}\n{traceback.format_exc()}")
raise
@spaces.GPU(duration=3)
@torch.no_grad()
def run_texture(image, mesh_path, seed, req=None):
try:
log_gpu_memory()
height, width = 512, 512
cameras = get_orthogonal_camera(
elevation_deg=[0, 0, 0, 89.99],
distance=[1.8] * NUM_VIEWS,
left=-0.55,
right=0.55,
bottom=-0.55,
top=0.55,
azimuth_deg=[x - 90 for x in [0, 90, 180, 180]],
device=DEVICE,
)
ctx = NVDiffRastContextWrapper(device=DEVICE, context_type="cuda")
mesh = load_mesh(mesh_path, rescale=True, device=DEVICE)
with autocast():
render_out = render(
ctx,
mesh,
cameras,
height=height,
width=width,
render_attr=False,
normal_background=0.0,
)
control_images = (
torch.cat(
[(render_out.pos + 0.5).clamp(0, 1), (render_out.normal / 2 + 0.5).clamp(0, 1)],
dim=-1,
)
.permute(0, 3, 1, 2)
.to(DEVICE)
)
del render_out
image = Image.open(image)
birefnet.to(DEVICE, dtype=DTYPE)
with autocast():
image = remove_bg_fn(image)
birefnet.to("cpu")
image = preprocess_image(image, height, width)
pipe_kwargs = {"generator": torch.Generator(device=DEVICE).manual_seed(seed)} if seed != -1 else {}
mv_adapter_pipe.to(DEVICE, dtype=DTYPE)
with autocast():
images = mv_adapter_pipe(
"high quality",
height=height,
width=width,
num_inference_steps=10,
guidance_scale=3.0,
num_images_per_prompt=NUM_VIEWS,
control_image=control_images,
control_conditioning_scale=1.0,
reference_image=image,
reference_conditioning_scale=1.0,
negative_prompt="watermark, ugly, deformed, noisy, blurry, low contrast",
cross_attention_kwargs={"scale": 1.0},
**pipe_kwargs,
).images
mv_adapter_pipe.to("cpu")
del control_images
save_dir = os.path.join(TMP_DIR, str(req.session_hash) if req else "examples")
os.makedirs(save_dir, exist_ok=True)
mv_image_path = os.path.join(save_dir, f"mv_adapter_{get_random_hex()}.png")
make_image_grid(images, rows=1).save(mv_image_path)
from texture import TexturePipeline, ModProcessConfig
texture_pipe = TexturePipeline(
upscaler_ckpt_path="checkpoints/RealESRGAN_x2plus.pth",
inpaint_ckpt_path="checkpoints/big-lama.pt",
device=DEVICE,
)
textured_glb_path = texture_pipe(
mesh_path=mesh_path,
save_dir=save_dir,
save_name=f"polygenixai_texture_mesh_{get_random_hex()}.glb",
uv_unwarp=True,
uv_size=2048,
rgb_path=mv_image_path,
rgb_process_config=ModProcessConfig(view_upscale=True, inpaint_mode="view"),
camera_azimuth_deg=[x - 90 for x in [0, 90, 180, 180]],
)
torch.cuda.empty_cache()
log_gpu_memory()
return textured_glb_path
except Exception as e:
logger.error(f"Error in run_texture: {str(e)}\n{traceback.format_exc()}")
raise
@spaces.GPU(duration=3)
@torch.no_grad()
def run_full(image, seed=0, num_inference_steps=30, guidance_scale=7.0, simplify=True, target_face_num=DEFAULT_FACE_NUMBER, req=None):
try:
log_gpu_memory()
image_seg = run_segmentation(image)
mesh_path = image_to_3d(image_seg, seed, num_inference_steps, guidance_scale, simplify, target_face_num, req)
textured_glb_path = run_texture(image, mesh_path, seed, req)
return image_seg, mesh_path, textured_glb_path
except Exception as e:
logger.error(f"Error in run_full: {str(e)}\n{traceback.format_exc()}")
raise
def gradio_generate(image, seed=0, num_inference_steps=30, guidance_scale=7.0, simplify=True, target_face_num=DEFAULT_FACE_NUMBER):
try:
logger.info("Starting gradio_generate")
api_key = os.getenv("POLYGENIX_API_KEY", "your-secret-api-key")
request = gr.Request()
if not request.headers.get("x-api-key") == api_key:
logger.error("Invalid API key")
raise ValueError("Invalid API key")
if image.startswith("data:image"):
logger.info("Processing base64 image")
base64_string = image.split(",")[1]
image_data = base64.b64decode(base64_string)
temp_image_path = os.path.join(TMP_DIR, f"input_{get_random_hex()}.png")
with open(temp_image_path, "wb") as f:
f.write(image_data)
else:
temp_image_path = image
if not os.path.exists(temp_image_path):
logger.error(f"Image file not found: {temp_image_path}")
raise ValueError("Invalid or missing image file")
image_seg, mesh_path, textured_glb_path = run_full(temp_image_path, seed, num_inference_steps, guidance_scale, simplify, target_face_num, request)
session_hash = os.path.basename(os.path.dirname(textured_glb_path))
logger.info(f"Generated model at /files/{session_hash}/{os.path.basename(textured_glb_path)}")
return {"file_url": f"/files/{session_hash}/{os.path.basename(textured_glb_path)}"}
except Exception as e:
logger.error(f"Error in gradio_generate: {str(e)}\n{traceback.format_exc()}")
raise
def start_session(req: gr.Request):
try:
save_dir = os.path.join(TMP_DIR, str(req.session_hash))
os.makedirs(save_dir, exist_ok=True)
logger.info(f"Started session, created directory: {save_dir}")
except Exception as e:
logger.error(f"Error in start_session: {str(e)}\n{traceback.format_exc()}")
raise
def end_session(req: gr.Request):
try:
save_dir = os.path.join(TMP_DIR, str(req.session_hash))
shutil.rmtree(save_dir)
logger.info(f"Ended session, removed directory: {save_dir}")
except Exception as e:
logger.error(f"Error in end_session: {str(e)}\n{traceback.format_exc()}")
raise
def get_random_seed(randomize_seed, seed):
try:
if randomize_seed:
seed = random.randint(0, MAX_SEED)
logger.info(f"Generated seed: {seed}")
return seed
except Exception as e:
logger.error(f"Error in get_random_seed: {str(e)}\n{traceback.format_exc()}")
raise
def download_image(url: str, save_path: str) -> str:
try:
logger.info(f"Downloading image from {url}")
response = requests.get(url, stream=True)
response.raise_for_status()
with open(save_path, "wb") as f:
for chunk in response.iter_content(chunk_size=8192):
f.write(chunk)
logger.info(f"Saved image to {save_path}")
return save_path
except Exception as e:
logger.error(f"Failed to download image from {url}: {str(e)}\n{traceback.format_exc()}")
raise
@spaces.GPU(duration=3)
@torch.no_grad()
def run_full_api(image, seed=0, num_inference_steps=30, guidance_scale=7.0, simplify=True, target_face_num=DEFAULT_FACE_NUMBER, req=None):
try:
logger.info("Running run_full_api")
def execute():
image_seg, mesh_path, textured_glb_path = run_full(image, seed, num_inference_steps, guidance_scale, simplify, target_face_num, req)
session_hash = os.path.basename(os.path.dirname(textured_glb_path))
return {"file_url": f"/files/{session_hash}/{os.path.basename(textured_glb_path)}"}
return retry_on_failure(execute)
except Exception as e:
logger.error(f"Error in run_full_api: {str(e)}\n{traceback.format_exc()}")
raise
# Define Gradio API endpoint
try:
logger.info("Initializing Gradio API interface")
api_interface = gr.Interface(
fn=gradio_generate,
inputs=[
gr.Image(type="filepath", label="Image"),
gr.Number(label="Seed", value=0, precision=0),
gr.Number(label="Inference Steps", value=30, precision=0),
gr.Number(label="Guidance Scale", value=7.0),
gr.Checkbox(label="Simplify Mesh", value=True),
gr.Number(label="Target Face Number", value=DEFAULT_FACE_NUMBER, precision=0)
],
outputs="json",
api_name="/api/generate"
)
logger.info("Gradio API interface initialized successfully")
except Exception as e:
logger.error(f"Failed to initialize Gradio API interface: {str(e)}\n{traceback.format_exc()}")
raise
HEADER = """
# 🌌 PolyGenixAI: Craft 3D Worlds with Cosmic Precision
## Unleash Infinite Creativity with AI-Powered 3D Generation by AnvilInteractive Solutions
<p style="font-size: 1.1em; color: #A78BFA;">By <a href="https://www.anvilinteractive.com/" style="color: #A78BFA; text-decoration: none; font-weight: bold;">AnvilInteractive Solutions</a></p>
## 🚀 Launch Your Creation:
1. **Upload an Image** (clear, single-object images shine brightest)
2. **Choose a Style Filter** to infuse your unique vision
3. Click **Generate 3D Model** to sculpt your mesh
4. Click **Apply Texture** to bring your model to life
5. **Download GLB** to share your masterpiece
<p style="font-size: 0.9em; margin-top: 10px; color: #D1D5DB;">Powered by cutting-edge AI and multi-view technology from AnvilInteractive Solutions. Join our <a href="https://www.anvilinteractive.com/community" style="color: #A78BFA; text-decoration: none;">PolyGenixAI Community</a> to connect with creators and spark inspiration.</p>
<style>
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;600;700&display=swap');
body {
background-color: #1A1A1A !important;
font-family: 'Inter', sans-serif !important;
color: #D1D5DB !important;
}
.gr-panel {
background-color: #2D2D2D !important;
border: 1px solid #7C3AED !important;
border-radius: 12px !important;
padding: 20px !important;
box-shadow: 0 4px 10px rgba(124, 58, 237, 0.2) !important;
}
.gr-button-primary {
background: linear-gradient(45deg, #7C3AED, #A78BFA) !important;
color: white !important;
border: none !important;
border-radius: 8px !important;
padding: 12px 24px !important;
font-weight: 600 !important;
transition: transform 0.2s, box-shadow 0.2s !important;
}
.gr-button-primary:hover {
transform: translateY(-2px) !important;
box-shadow: 0 4px 12px rgba(124, 58, 237, 0.5) !important;
}
.gr-button-secondary {
background-color: #4B4B4B !important;
color: #D1D5DB !important;
border: 1px solid #A78BFA !important;
border-radius: 8px !important;
padding: 10px 20px !important;
transition: transform 0.2s !important;
}
.gr-button-secondary:hover {
transform: translateY(-1px) !important;
background-color: #6B6B6B !important;
}
.gr-accordion {
background-color: #2D2D2D !important;
border-radius: 8px !important;
border: 1px solid #7C3AED !important;
}
.gr-tab {
background-color: #2D2D2D !important;
color: #A78BFA !important;
border: 1px solid #7C3AED !important;
border-radius: 8px !important;
margin: 5px !important;
}
.gr-tab:hover, .gr-tab-selected {
background: linear-gradient(45deg, #7C3AED, #A78BFA) !important;
color: white !important;
}
.gr-slider input[type=range]::-webkit-slider-thumb {
background-color: #7C3AED !important;
border: 2px solid #A78BFA !important;
}
.gr-dropdown {
background-color: #2D2D2D !important;
color: #D1D5DB !important;
border: 1px solid #A78BFA !important;
border-radius: 8px !important;
}
h1, h3 {
color: #A78BFA !important;
text-shadow: 0 0 10px rgba(124, 58, 237, 0.5) !important;
}
</style>
"""
try:
logger.info("Initializing Gradio Blocks interface")
with gr.Blocks(title="PolyGenixAI", css="body { background-color: #1A1A1A; } .gr-panel { background-color: #2D2D2D; }") as demo:
gr.Markdown(HEADER)
with gr.Tabs(elem_classes="gr-tab"):
with gr.Tab("Create 3D Model"):
with gr.Row():
with gr.Column(scale=1):
image_prompts = gr.Image(label="Upload Image", type="filepath", height=300, elem_classes="gr-panel")
seg_image = gr.Image(label="Preview Segmentation", type="pil", format="png", interactive=False, height=300, elem_classes="gr-panel")
with gr.Accordion("Style & Settings", open=True, elem_classes="gr-accordion"):
style_filter = gr.Dropdown(
choices=["None", "Realistic", "Fantasy", "Cartoon", "Sci-Fi", "Vintage", "Cosmic", "Neon"],
label="Style Filter",
value="None",
info="Select a style to inspire your 3D model (optional)",
elem_classes="gr-dropdown"
)
seed = gr.Slider(
label="Seed",
minimum=0,
maximum=MAX_SEED,
step=1,
value=0,
elem_classes="gr-slider"
)
randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
num_inference_steps = gr.Slider(
label="Inference Steps",
minimum=8,
maximum=50,
step=1,
value=30,
info="Higher steps enhance detail but increase processing time",
elem_classes="gr-slider"
)
guidance_scale = gr.Slider(
label="Guidance Scale",
minimum=0.0,
maximum=20.0,
step=0.1,
value=7.0,
info="Controls adherence to input image",
elem_classes="gr-slider"
)
reduce_face = gr.Checkbox(label="Simplify Mesh", value=True)
target_face_num = gr.Slider(
maximum=100000,
minimum=10000,
value=DEFAULT_FACE_NUMBER,
label="Target Face Number",
info="Adjust mesh complexity for performance",
elem_classes="gr-slider"
)
gen_button = gr.Button("Generate 3D Model", variant="primary", elem_classes="gr-button-primary")
gen_texture_button = gr.Button("Apply Texture", variant="secondary", interactive=False, elem_classes="gr-button-secondary")
with gr.Column(scale=1):
model_output = gr.Model3D(label="3D Model Preview", interactive=False, height=400, elem_classes="gr-panel")
textured_model_output = gr.Model3D(label="Textured 3D Model", interactive=False, height=400, elem_classes="gr-panel")
download_button = gr.Button("Download GLB", variant="secondary", elem_classes="gr-button-secondary")
with gr.Tab("Cosmic Gallery"):
gr.Markdown("### Discover Stellar Creations")
gr.Examples(
examples=[
f"{TRIPOSG_CODE_DIR}/assets/example_data/{image}"
for image in os.listdir(f"{TRIPOSG_CODE_DIR}/assets/example_data")
],
fn=run_full_api,
inputs=[image_prompts],
outputs=[seg_image, model_output, textured_model_output],
cache_examples=True,
)
gr.Markdown("Connect with creators in our <a href='https://www.anvilinteractive.com/community' style='color: #A78BFA; text-decoration: none;'>PolyGenixAI Cosmic Community</a>!")
gen_button.click(
run_segmentation,
inputs=[image_prompts],
outputs=[seg_image]
).then(
get_random_seed,
inputs=[randomize_seed, seed],
outputs=[seed],
).then(
image_to_3d,
inputs=[
seg_image,
seed,
num_inference_steps,
guidance_scale,
reduce_face,
target_face_num
],
outputs=[model_output]
).then(
lambda: gr.Button(interactive=True), outputs=[gen_texture_button]
)
gen_texture_button.click(
run_texture,
inputs=[image_prompts, model_output, seed],
outputs=[textured_model_output]
)
demo.load(start_session)
demo.unload(end_session)
logger.info("Gradio Blocks interface initialized successfully")
except Exception as e:
logger.error(f"Failed to initialize Gradio Blocks interface: {str(e)}\n{traceback.format_exc()}")
raise
if __name__ == "__main__":
try:
logger.info("Launching Gradio application")
demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)
logger.info("Gradio application launched successfully")
except Exception as e:
logger.error(f"Failed to launch Gradio application: {str(e)}\n{traceback.format_exc()}")
raise