import spaces import logging import os import random import re import sys import tempfile import uuid import warnings import atexit import threading from pathlib import Path from io import BytesIO import numpy as np from PIL import Image from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler import gradio as gr import torch from transformers import AutoModelForCausalLM, AutoTokenizer from plyfile import PlyData sys.path.append(os.path.dirname(os.path.abspath(__file__))) sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), "ml-sharp", "src")) from diffusers import ZImagePipeline from diffusers.models.transformers.transformer_z_image import ZImageTransformer2DModel # SHARP imports for 3D Gaussian splat generation from sharp.models import create_predictor, PredictorParams from sharp.utils.gaussians import save_ply from sharp.cli.predict import predict_image, DEFAULT_MODEL_URL # ==================== Environment Variables ================================== MODEL_PATH = os.environ.get("MODEL_PATH", "Tongyi-MAI/Z-Image-Turbo") ENABLE_COMPILE = os.environ.get("ENABLE_COMPILE", "false").lower() == "true" ENABLE_WARMUP = os.environ.get("ENABLE_WARMUP", "false").lower() == "true" ATTENTION_BACKEND = os.environ.get("ATTENTION_BACKEND", "native") HF_TOKEN = os.environ.get("HF_TOKEN") # ============================================================================= os.environ["TOKENIZERS_PARALLELISM"] = "false" warnings.filterwarnings("ignore") logging.getLogger("transformers").setLevel(logging.ERROR) # Temporary file cleanup system _temp_files_lock = threading.Lock() _temp_files = [] def register_temp_file(path: str): """Register a temporary file for cleanup.""" with _temp_files_lock: _temp_files.append(path) def cleanup_temp_files(): """Clean up all registered temporary files.""" with _temp_files_lock: for path in _temp_files: try: if os.path.exists(path): os.unlink(path) except Exception as e: print(f"Failed to delete temp file {path}: {e}") _temp_files.clear() atexit.register(cleanup_temp_files) def optimize_memory(): """Clear CUDA cache and run garbage collection.""" import gc gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() torch.cuda.synchronize() RES_CHOICES = { "1024": [ "1024x1024 ( 1:1 )", "1152x896 ( 9:7 )", "896x1152 ( 7:9 )", "1152x864 ( 4:3 )", "864x1152 ( 3:4 )", "1248x832 ( 3:2 )", "832x1248 ( 2:3 )", "1280x720 ( 16:9 )", "720x1280 ( 9:16 )", "1344x576 ( 21:9 )", "576x1344 ( 9:21 )", ], "1280": [ "1280x1280 ( 1:1 )", "1440x1120 ( 9:7 )", "1120x1440 ( 7:9 )", "1472x1104 ( 4:3 )", "1104x1472 ( 3:4 )", "1536x1024 ( 3:2 )", "1024x1536 ( 2:3 )", "1536x864 ( 16:9 )", "864x1536 ( 9:16 )", "1680x720 ( 21:9 )", "720x1680 ( 9:21 )", ], "1536": [ "1536x1536 ( 1:1 )", "1728x1344 ( 9:7 )", "1344x1728 ( 7:9 )", "1728x1296 ( 4:3 )", "1296x1728 ( 3:4 )", "1872x1248 ( 3:2 )", "1248x1872 ( 2:3 )", "2048x1152 ( 16:9 )", "1152x2048 ( 9:16 )", "2016x864 ( 21:9 )", "864x2016 ( 9:21 )", ], } RESOLUTION_SET = [] for resolutions in RES_CHOICES.values(): RESOLUTION_SET.extend(resolutions) EXAMPLE_PROMPTS = [ ["Alien UFO landing in a dark forest with a starry sky"], [ "Underwater city with futuristic buildings and colorful coral reefs, vibrant marine life swimming around, sunlight filtering through the water, digital art" ], [ "A serene mountain landscape during autumn, with a clear blue lake reflecting the colorful foliage, high-resolution photograph" ], [ "A bustling cyberpunk city street at night, neon signs in various languages,style of Syd Mead and Katsuhiro Otomo" ], ] def get_resolution(resolution: str) -> tuple[int, int]: """Parse resolution string to width and height tuple.""" match = re.search(r"(\d+)\s*[×x]\s*(\d+)", resolution) if match: return int(match.group(1)), int(match.group(2)) return 1024, 1024 def load_models(model_path: str, enable_compile: bool = False, attention_backend: str = "native"): """ Load all models required for Z-Image generation. Uses device_map="cuda" for ZeroGPU compatibility. """ print(f"Loading models from {model_path}...") use_auth_token = HF_TOKEN if HF_TOKEN else True is_local = os.path.exists(model_path) # Load VAE if is_local: vae = AutoencoderKL.from_pretrained( os.path.join(model_path, "vae"), torch_dtype=torch.bfloat16, device_map="cuda", ) else: vae = AutoencoderKL.from_pretrained( model_path, subfolder="vae", torch_dtype=torch.bfloat16, device_map="cuda", use_auth_token=use_auth_token, ) # Load Text Encoder if is_local: text_encoder = AutoModelForCausalLM.from_pretrained( os.path.join(model_path, "text_encoder"), torch_dtype=torch.bfloat16, device_map="cuda", ).eval() else: text_encoder = AutoModelForCausalLM.from_pretrained( model_path, subfolder="text_encoder", torch_dtype=torch.bfloat16, device_map="cuda", use_auth_token=use_auth_token, ).eval() # Load Tokenizer if is_local: tokenizer = AutoTokenizer.from_pretrained(os.path.join(model_path, "tokenizer")) else: tokenizer = AutoTokenizer.from_pretrained( model_path, subfolder="tokenizer", use_auth_token=use_auth_token, ) tokenizer.padding_side = "left" # Configure torch.compile optimizations if enable_compile: print("Enabling torch.compile optimizations...") torch._inductor.config.conv_1x1_as_mm = True torch._inductor.config.coordinate_descent_tuning = True torch._inductor.config.epilogue_fusion = False torch._inductor.config.coordinate_descent_check_all_directions = True torch._inductor.config.max_autotune_gemm = True torch._inductor.config.max_autotune_gemm_backends = "TRITON,ATEN" torch._inductor.config.triton.cudagraphs = False # Create pipeline pipe = ZImagePipeline( scheduler=None, vae=vae, text_encoder=text_encoder, tokenizer=tokenizer, transformer=None, ) if enable_compile: pipe.vae.disable_tiling() # Load Transformer if is_local: transformer = ZImageTransformer2DModel.from_pretrained( os.path.join(model_path, "transformer") ).to("cuda", torch.bfloat16) else: transformer = ZImageTransformer2DModel.from_pretrained( model_path, subfolder="transformer", use_auth_token=use_auth_token, ).to("cuda", torch.bfloat16) pipe.transformer = transformer pipe.transformer.set_attention_backend(attention_backend) if enable_compile: print("Compiling transformer...") pipe.transformer = torch.compile( pipe.transformer, mode="max-autotune-no-cudagraphs", fullgraph=False ) pipe.to("cuda", torch.bfloat16) print("Models loaded successfully") return pipe def generate_image( pipe, prompt: str, resolution: str = "1024x1024", seed: int = 42, guidance_scale: float = 0.0, num_inference_steps: int = 9, shift: float = 3.0, max_sequence_length: int = 512, progress=gr.Progress(track_tqdm=True), ): """Generate a single image using the Z-Image pipeline.""" width, height = get_resolution(resolution) generator = torch.Generator("cuda").manual_seed(seed) scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=1000, shift=shift) pipe.scheduler = scheduler image = pipe( prompt=prompt, height=height, width=width, guidance_scale=guidance_scale, num_inference_steps=num_inference_steps, generator=generator, max_sequence_length=max_sequence_length, ).images[0] return image def warmup_model(pipe, resolutions: list[str]): """Warm up the model by running inference on dummy prompts.""" print("Starting warmup phase...") dummy_prompt = "warmup" for res_str in resolutions: print(f"Warming up for resolution: {res_str}") try: for i in range(3): generate_image( pipe, prompt=dummy_prompt, resolution=res_str, num_inference_steps=9, guidance_scale=0.0, seed=42 + i, ) except Exception as e: print(f"Warmup failed for {res_str}: {e}") print("Warmup completed.") # Global pipeline instance pipe = None # Global SHARP splat predictor (lazy loaded) splat_predictor = None def init_app(): """Initialize the application by loading models.""" global pipe try: pipe = load_models( MODEL_PATH, enable_compile=ENABLE_COMPILE, attention_backend=ATTENTION_BACKEND, ) print(f"Model loaded. Compile: {ENABLE_COMPILE}, Backend: {ATTENTION_BACKEND}") if ENABLE_WARMUP: all_resolutions = [] for cat in RES_CHOICES.values(): all_resolutions.extend(cat) warmup_model(pipe, all_resolutions) except Exception as e: print(f"Error loading model: {e}") import traceback traceback.print_exc() pipe = None def load_splat_predictor(device: str = "cuda"): """Load the SHARP Gaussian splat predictor model (lazy loading).""" global splat_predictor if splat_predictor is not None: return splat_predictor print(f"Loading SHARP splat predictor from {DEFAULT_MODEL_URL}...") # Download and load the model state_dict = torch.hub.load_state_dict_from_url( DEFAULT_MODEL_URL, progress=True, ) predictor = create_predictor(PredictorParams()) predictor.load_state_dict(state_dict) predictor.eval() predictor.to(device) splat_predictor = predictor print("SHARP predictor loaded successfully") return splat_predictor def convert_ply_to_splat(ply_file_path: str) -> bytes: """ Convert a PLY file to SPLAT format for the antimatter15 viewer. Returns the splat data as bytes. """ plydata = PlyData.read(ply_file_path) vert = plydata["vertex"] sorted_indices = np.argsort( -np.exp(vert["scale_0"] + vert["scale_1"] + vert["scale_2"]) / (1 + np.exp(-vert["opacity"])) ) buffer = BytesIO() for idx in sorted_indices: v = plydata["vertex"][idx] position = np.array([v["x"], v["y"], v["z"]], dtype=np.float32) scales = np.exp( np.array([v["scale_0"], v["scale_1"], v["scale_2"]], dtype=np.float32) ) color = np.array([ 0.5 + 0.28209479177387814 * v["f_dc_0"], 0.5 + 0.28209479177387814 * v["f_dc_1"], 0.5 + 0.28209479177387814 * v["f_dc_2"], 1 / (1 + np.exp(-v["opacity"])), ]) rot = np.array([v["rot_0"], v["rot_1"], v["rot_2"], v["rot_3"]], dtype=np.float32) buffer.write(position.tobytes()) buffer.write(scales.tobytes()) buffer.write((color * 255).clip(0, 255).astype(np.uint8).tobytes()) buffer.write( ((rot / np.linalg.norm(rot)) * 128 + 128).clip(0, 255).astype(np.uint8).tobytes() ) return buffer.getvalue() @spaces.GPU def generate_splat(selected_image, progress=gr.Progress(track_tqdm=True)): """Generate a 3D Gaussian splat from the selected image.""" if selected_image is None: raise gr.Error("Please select an image from the gallery first") try: if isinstance(selected_image, str): if not os.path.exists(selected_image): raise gr.Error(f"Image file not found: {selected_image}") pil_image = Image.open(selected_image).convert("RGB") image_np = np.array(pil_image) elif hasattr(selected_image, "convert"): image_np = np.array(selected_image.convert("RGB")) elif isinstance(selected_image, np.ndarray): image_np = selected_image.copy() else: image_np = np.array(selected_image) if image_np is None or image_np.size == 0: raise gr.Error("Invalid image data") if image_np.ndim == 2: image_np = np.stack([image_np] * 3, axis=-1) elif image_np.ndim == 3 and image_np.shape[-1] == 4: image_np = image_np[:, :, :3] elif image_np.ndim != 3 or image_np.shape[-1] != 3: raise gr.Error(f"Unexpected image shape: {image_np.shape}") height, width = image_np.shape[:2] if height < 64 or width < 64: raise gr.Error(f"Image too small: {width}x{height}. Minimum is 64x64.") f_mm = 30.0 f_px = f_mm * np.sqrt(width**2 + height**2) / np.sqrt(36**2 + 24**2) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") predictor = load_splat_predictor(str(device)) print(f"Generating 3D Gaussian splat from {width}x{height} image...") gaussians = predict_image(predictor, image_np, f_px, device) num_gaussians = gaussians.mean_vectors.shape[1] print(f"Generated {num_gaussians} gaussians") # Use gradio's temp directory for better compatibility with Spaces cache_dir = Path(tempfile.gettempdir()) / "gradio_cache" cache_dir.mkdir(exist_ok=True) ply_path = cache_dir / f"scene_{uuid.uuid4().hex[:8]}.ply" save_ply(gaussians, f_px, (height, width), ply_path) register_temp_file(str(ply_path)) status_msg = f"✅ Generated {num_gaussians:,} gaussians | PLY: {ply_path.stat().st_size/1024:.1f}KB" return ( str(ply_path), # ply_download - return path directly str(ply_path), # ply_path_state gr.update(visible=True), # convert_splat_btn status_msg, # splat_status ) except gr.Error: raise except Exception as e: print(f"Error generating splat: {e}") import traceback traceback.print_exc() raise gr.Error(f"Failed to generate 3D splat: {str(e)}") def convert_and_save_splat(ply_path): """Convert PLY to SPLAT format and save to temp file.""" if not ply_path or not os.path.exists(ply_path): raise gr.Error("PLY file not found. Please generate a 3D splat first.") try: splat_data = convert_ply_to_splat(ply_path) # Use same cache directory as PLY files cache_dir = Path(tempfile.gettempdir()) / "gradio_cache" cache_dir.mkdir(exist_ok=True) splat_path = cache_dir / f"scene_{uuid.uuid4().hex[:8]}.splat" with open(splat_path, "wb") as f: f.write(splat_data) register_temp_file(str(splat_path)) size_kb = splat_path.stat().st_size / 1024 status_msg = f"✅ SPLAT file created | Size: {size_kb:.1f}KB" return status_msg, str(splat_path) # Return path directly except Exception as e: print(f"Error converting to SPLAT: {e}") raise gr.Error(f"Failed to convert to SPLAT: {str(e)}") @spaces.GPU def generate( prompt: str, resolution: str = "1024x1024 ( 1:1 )", seed: int = 42, steps: int = 9, shift: float = 3.0, random_seed: bool = True, gallery_images: list = None, progress=gr.Progress(track_tqdm=True), ): """ Generate an image using the Z-Image model based on the provided prompt and settings. Args: prompt: Text prompt describing the desired image content resolution: Output resolution in format "WIDTHxHEIGHT ( RATIO )" seed: Seed for reproducible generation steps: Number of inference steps for the diffusion process shift: Time shift parameter for the flow matching scheduler random_seed: Whether to generate a new random seed gallery_images: List of previously generated images to append to progress: Gradio progress tracker Returns: tuple: (gallery_images, seed_str, seed_int) """ if random_seed: new_seed = random.randint(1, 1000000) else: new_seed = seed if seed != -1 else random.randint(1, 1000000) if pipe is None: raise gr.Error("Model not loaded. Please check the console for errors.") # Parse resolution try: resolution_str = resolution.split(" ")[0] except: resolution_str = "1024x1024" # Generate image image = generate_image( pipe=pipe, prompt=prompt, resolution=resolution_str, seed=new_seed, guidance_scale=0.0, num_inference_steps=int(steps + 1), shift=shift, ) if gallery_images is None: gallery_images = [] # Add latest output to the top of the list gallery_images = [image] + gallery_images return gallery_images, new_seed @spaces.GPU def generate_batch( prompt: str, resolution: str, seed: int, steps: int, shift: float, batch_size: int = 2, gallery_images: list = None, progress=gr.Progress(track_tqdm=True), ): """Generate multiple images in a batch for efficiency.""" if pipe is None: raise gr.Error("Model not loaded.") if gallery_images is None: gallery_images = [] new_images = [] for i in range(batch_size): current_seed = seed + i image = generate_image( pipe=pipe, prompt=prompt, resolution=resolution.split(" ")[0], seed=current_seed, guidance_scale=0.0, num_inference_steps=int(steps + 1), shift=shift, ) new_images.append(image) optimize_memory() return new_images + gallery_images, seed # Initialize the app init_app() # ==================== Gradio UI ==================== css = """ .fillable{max-width: 1230px !important} """ with gr.Blocks(title="Z-Image Demo") as demo: gr.Markdown( """
# Generative 3D Gaussian Splat * Generate images from text prompts using [![GitHub](https://img.shields.io/badge/GitHub-Z--Image-181717?logo=github&logoColor=white)](https://github.com/Tongyi-MAI/Z-Image) * Create 3D Gaussian splat models from generated images using [![GitHub](https://img.shields.io/badge/GitHub-SHARP-181717?logo=github&logoColor=white)](https://github.com/apple/ml-sharp)
""" ) with gr.Row(): with gr.Column(scale=1): prompt_input = gr.Textbox( label="Prompt", lines=3, placeholder="Enter your prompt here...", ) with gr.Row(): choices = [int(k) for k in RES_CHOICES.keys()] res_cat = gr.Dropdown( value=1024, choices=choices, label="Resolution Category", ) initial_res_choices = RES_CHOICES["1024"] resolution = gr.Dropdown( value=initial_res_choices[0], choices=RESOLUTION_SET, label="Width x Height (Ratio)", ) with gr.Row(): seed = gr.Number(label="Seed", value=42, precision=0) random_seed = gr.Checkbox(label="Random Seed", value=True) with gr.Row(): steps = gr.Slider( label="Steps", minimum=1, maximum=100, value=8, step=1, ) shift = gr.Slider( label="Time Shift", minimum=1.0, maximum=10.0, value=3.0, step=0.1, ) generate_btn = gr.Button("Generate", variant="primary") # Example prompts gr.Markdown("### 📝 Example Prompts") gr.Examples(examples=EXAMPLE_PROMPTS, inputs=prompt_input, label=None) with gr.Column(scale=1): output_gallery = gr.Gallery( label="Generated Images", columns=2, rows=2, height=600, object_fit="contain", format="png", interactive=False, ) # 3D Gaussian Splat Generation Section with gr.Accordion("3D Gaussian Splat Generation", open=False): # gr.Markdown( # """Click on an image in the gallery above to select it, then click "Generate 3D Splat" to create # a 3D Gaussian splat model. You can then download PLY files or SPLAT in an additional step.""" # ) # State to hold selected image and PLY path selected_image_state = gr.State(value=None) ply_path_state = gr.State(value=None) with gr.Row(): generate_splat_btn = gr.Button( "Generate 3D Scene", variant="secondary", interactive=False, ) convert_splat_btn = gr.Button( "Convert to SPLAT", variant="secondary", visible=False, ) splat_status = gr.Textbox( label="Status", interactive=False, visible=True, value="Click an image in the gallery to select it", lines=2 ) # Download section - always visible, empty until files are generated with gr.Row(): ply_download = gr.File( label="PLY File (click to download)", visible=True, interactive=False, value=None, ) splat_download = gr.File( label="SPLAT File (click to download)", visible=True, interactive=False, value=None, ) # Gallery image selection handler def on_gallery_select(evt: gr.SelectData, gallery_images): """Handle gallery image selection.""" if gallery_images is None or len(gallery_images) == 0: return None, gr.update(interactive=False), "No image selected" selected_idx = evt.index if selected_idx < len(gallery_images): selected_item = gallery_images[selected_idx] if isinstance(selected_item, tuple): selected_img = selected_item[0] elif isinstance(selected_item, dict): selected_img = selected_item.get('image') or selected_item.get('name') else: selected_img = selected_item return ( selected_img, gr.update(interactive=True), f"Image {selected_idx + 1} selected - ready to generate 3D splat", ) return None, gr.update(interactive=False), "Selection error" output_gallery.select( on_gallery_select, inputs=[output_gallery], outputs=[selected_image_state, generate_splat_btn, splat_status], ) # Splat generation handler generate_splat_btn.click( generate_splat, inputs=[selected_image_state], outputs=[ply_download, ply_path_state, convert_splat_btn, splat_status], ) # SPLAT conversion handler convert_splat_btn.click( convert_and_save_splat, inputs=[ply_path_state], outputs=[splat_status, splat_download], ) def update_res_choices(res_cat_value): """Update resolution choices based on selected category.""" if str(res_cat_value) in RES_CHOICES: res_choices = RES_CHOICES[str(res_cat_value)] else: res_choices = RES_CHOICES["1024"] return gr.update(value=res_choices[0], choices=res_choices) res_cat.change( update_res_choices, inputs=res_cat, outputs=resolution, ) generate_btn.click( generate, inputs=[prompt_input, resolution, seed, steps, shift, random_seed, output_gallery], outputs=[output_gallery, seed], ) if __name__ == "__main__": demo.launch(css=css)