import gradio as gr import spaces from gradio_litmodel3d import LitModel3D import numpy as np import random import os import torch from diffusers import DiffusionPipeline from transformers import pipeline, AutoTokenizer from huggingface_hub import login from PIL import Image import shutil os.environ['SPCONV_ALGO'] = 'native' from typing import * import imageio from easydict import EasyDict as edict from trellis.pipelines import TrellisImageTo3DPipeline from trellis.representations import Gaussian, MeshExtractResult from trellis.utils import render_utils, postprocessing_utils hf_token = os.getenv("hf_token") login(token=hf_token) # Global constants and default values MAX_SEED = np.iinfo(np.int32).max MAX_IMAGE_SIZE = 2048 # Default system prompt for text generation DEFAULT_SYSTEM_PROMPT = """You are a product designer with strong knowledge in text-to-image generation. You will receive a product request in the form of a brief description, and your mission will be to imagine a new product design that meets this need. The deliverable (generated response) will be exclusively a text prompt for the FLUX.1-dev text-to-image AI. This prompt should include a visual description of the object explicitly mentioning the essential aspects of its function. Additionally, you should explicitly mention in this prompt the aesthetic/photo characteristics of the image rendering (e.g., photorealistic, high quality, focal length, grain, etc.), knowing that the image will be the main image of this object in the product catalog. The background of the generated image must be entirely white. The prompt should be without narration.""" # Default Flux parameters DEFAULT_SEED = 42 DEFAULT_RANDOMIZE_SEED = True DEFAULT_WIDTH = 512 DEFAULT_HEIGHT = 512 DEFAULT_NUM_INFERENCE_STEPS = 6 DEFAULT_GUIDANCE_SCALE = 0.0 DEFAULT_TEMPERATURE = 0.9 TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tmp') os.makedirs(TMP_DIR, exist_ok=True) _text_gen_pipeline = None _image_gen_pipeline = None _trellis_pipeline = None def start_session(req: gr.Request): # user_dir = os.path.join(TMP_DIR, "temp_output") user_dir = os.path.join(TMP_DIR, str(req.session_hash)) os.makedirs(user_dir, exist_ok=True) def end_session(req: gr.Request): # user_dir = os.path.join(TMP_DIR, "temp_output") user_dir = os.path.join(TMP_DIR, str(req.session_hash)) shutil.rmtree(user_dir) def preprocess_image(image: Image.Image) -> Image.Image: trellis = get_trellis_pipeline() if trellis is None: # If the pipeline is not loaded, just return the original image return image processed_image = trellis.preprocess_image(image) return processed_image @spaces.GPU() def get_image_gen_pipeline(): global _image_gen_pipeline if (_image_gen_pipeline is None): try: device = "cuda" if torch.cuda.is_available() else "cpu" dtype = torch.bfloat16 _image_gen_pipeline = DiffusionPipeline.from_pretrained( # "black-forest-labs/FLUX.1-schnell", "black-forest-labs/FLUX.1-dev", torch_dtype=dtype, ).to(device) except Exception as e: print(f"Error loading image generation model: {e}") return None return _image_gen_pipeline @spaces.GPU() def get_text_gen_pipeline(): global _text_gen_pipeline if (_text_gen_pipeline is None): try: device = "cuda" if torch.cuda.is_available() else "cpu" tokenizer = AutoTokenizer.from_pretrained( # "mistralai/Mistral-7B-Instruct-v0.3", "deepseek-ai/DeepSeek-R1-Distill-Llama-8B", use_fast=True ) tokenizer.pad_token = tokenizer.pad_token or tokenizer.eos_token _text_gen_pipeline = pipeline( # "text-generation", # model="mistralai/Mistral-7B-Instruct-v0.3", model="deepseek-ai/DeepSeek-R1-Distill-Llama-8B", tokenizer=tokenizer, max_new_tokens=2048, device=device, pad_token_id=tokenizer.pad_token_id ) except Exception as e: print(f"Error loading text generation model: {e}") return None return _text_gen_pipeline # @spaces.GPU() def get_trellis_pipeline(): global _trellis_pipeline if _trellis_pipeline is None: try: print("Loading Trellis pipeline...") _trellis_pipeline = TrellisImageTo3DPipeline.from_pretrained("JeffreyXiang/TRELLIS-image-large") except Exception as e: print(f"Error loading Trellis pipeline: {e}") return None return _trellis_pipeline @spaces.GPU() def refine_prompt( prompt, system_prompt=DEFAULT_SYSTEM_PROMPT, progress=gr.Progress(track_tqdm=True) ): text_gen = get_text_gen_pipeline() if text_gen is None: return "", "Text generation model is unavailable." try: messages = [ {"role": "system", "content": system_prompt}, {"role": "user", "content": prompt}, ] # Indicate progress started progress(0, desc="Generating text") # Generate text refined_prompt = text_gen(messages) # Indicate progress complete progress(1) # Extract just the assistant's content from the response try: messages = refined_prompt[0]['generated_text'] # Find the last message with role 'assistant' assistant_messages = [msg for msg in messages if msg['role'] == 'assistant'] if not assistant_messages: return "", "Error: No assistant response found" assistant_content = assistant_messages[-1]['content'] # Remove quotation marks at the beginning and end if assistant_content.startswith('"') and assistant_content.endswith('"'): assistant_content = assistant_content[1:-1] return assistant_content, "Prompt refined successfully!" except (KeyError, IndexError): return "", "Error: Unexpected response format from the model" except Exception as e: print(f"Error in refine_prompt: {str(e)}") # Add debug print return "", f"Error refining prompt: {str(e)}" def validate_dimensions(width, height): if width * height > MAX_IMAGE_SIZE * MAX_IMAGE_SIZE: return False, "Image dimensions too large" return True, None @spaces.GPU() def generate_image(prompt, seed=DEFAULT_SEED, randomize_seed=DEFAULT_RANDOMIZE_SEED, width=DEFAULT_WIDTH, height=DEFAULT_HEIGHT, num_inference_steps=DEFAULT_NUM_INFERENCE_STEPS, progress=gr.Progress(track_tqdm=True)): try: # Validate that prompt is not empty if not prompt or prompt.strip() == "": return None, "Please provide a valid prompt." progress(0.1, desc="Loading model") pipe = get_image_gen_pipeline() if pipe is None: return None, "Image generation model is unavailable." is_valid, error_msg = validate_dimensions(width, height) if not is_valid: return None, error_msg if randomize_seed: seed = random.randint(0, MAX_SEED) # Use default torch generator instead of cuda-specific generator generator = torch.Generator().manual_seed(seed) progress(0.3, desc="Running inference") # Match the working example's parameters output = pipe( prompt=prompt, width=width, height=height, num_inference_steps=num_inference_steps, generator=generator, guidance_scale=DEFAULT_GUIDANCE_SCALE, ) progress(0.8, desc="Processing output") image = output.images[0] progress(1.0, desc="Complete") return image, f"Image generated successfully with seed {seed}" except Exception as e: print(f"Error in generate_image: {str(e)}") return None, f"Error generating image: {str(e)}" examples = [ "a backpack for kids, flower style", "medieval flip flops", "cat shaped cake mold", ] css=""" #col-container { margin: 0 auto; max-width: 520px; } """ def pack_state(gs: Gaussian, mesh: MeshExtractResult) -> dict: return { 'gaussian': { **gs.init_params, '_xyz': gs._xyz.cpu().numpy(), '_features_dc': gs._features_dc.cpu().numpy(), '_scaling': gs._scaling.cpu().numpy(), '_rotation': gs._rotation.cpu().numpy(), '_opacity': gs._opacity.cpu().numpy(), }, 'mesh': { 'vertices': mesh.vertices.cpu().numpy(), 'faces': mesh.faces.cpu().numpy(), }, } def unpack_state(state: dict) -> Tuple[Gaussian, edict, str]: gs = Gaussian( aabb=state['gaussian']['aabb'], sh_degree=state['gaussian']['sh_degree'], mininum_kernel_size=state['gaussian']['mininum_kernel_size'], scaling_bias=state['gaussian']['scaling_bias'], opacity_bias=state['gaussian']['opacity_bias'], scaling_activation=state['gaussian']['scaling_activation'], ) gs._xyz = torch.tensor(state['gaussian']['_xyz'], device='cuda') gs._features_dc = torch.tensor(state['gaussian']['_features_dc'], device='cuda') gs._scaling = torch.tensor(state['gaussian']['_scaling'], device='cuda') gs._rotation = torch.tensor(state['gaussian']['_rotation'], device='cuda') gs._opacity = torch.tensor(state['gaussian']['_opacity'], device='cuda') mesh = edict( vertices=torch.tensor(state['mesh']['vertices'], device='cuda'), faces=torch.tensor(state['mesh']['faces'], device='cuda'), ) return gs, mesh @spaces.GPU def image_to_3d( image: Image.Image, seed: int, ss_guidance_strength: float, ss_sampling_steps: int, slat_guidance_strength: float, slat_sampling_steps: int, req: gr.Request, ) -> Tuple[dict, str]: try: # Load the Trellis pipeline pipeline = get_trellis_pipeline() if pipeline is None: return None, "Trellis pipeline is unavailable." pipeline.cuda() # Preprocess image image = preprocess_image(image) # Run the pipeline outputs = pipeline.run( image, seed=seed, formats=["gaussian", "mesh"], preprocess_image=False, sparse_structure_sampler_params={ "steps": ss_sampling_steps, "cfg_strength": ss_guidance_strength, }, slat_sampler_params={ "steps": slat_sampling_steps, "cfg_strength": slat_guidance_strength, }, ) # temp_dir = os.path.join(TMP_DIR, "temp_output") temp_dir = os.path.join(TMP_DIR, str(req.session_hash)) video = render_utils.render_video(outputs['gaussian'][0], num_frames=120)['color'] video_geo = render_utils.render_video(outputs['mesh'][0], num_frames=120)['normal'] video = [np.concatenate([video[i], video_geo[i]], axis=1) for i in range(len(video))] video_path = os.path.join(temp_dir, 'sample.mp4') imageio.mimsave(video_path, video, fps=15) state = pack_state(outputs['gaussian'][0], outputs['mesh'][0]) torch.cuda.empty_cache() return state, video_path except Exception as e: print(f"Error in image_to_3d: {str(e)}") import traceback traceback.print_exc() # Print the full stack trace for debugging return None, f"Error generating 3D model: {str(e)}" def process_example_pipeline(example_prompt): return example_prompt def create_interface(): model_status = "ℹ️ Models will be loaded on demand" with gr.Blocks(css=css) as demo: # Move session handlers INSIDE the Blocks context demo.load(fn=start_session) demo.unload(fn=end_session) gr.Info(model_status) # State for storing 3D model data output_state = gr.State(None) with gr.Column(elem_id="col-container"): gr.Markdown("# Text to 3D\nFocusing on product creation\nUsing Mistral-7B + FLUX.1-dev + Trellis") prompt = gr.Text( show_label=False, max_lines=1, placeholder="Enter basic object prompt", container=False, ) prompt_button = gr.Button("Refine prompt with Mistral") refined_prompt = gr.Text( show_label=False, max_lines=10, placeholder="Detailed object prompt", container=False, max_length=2048, ) visual_button = gr.Button("Create visual with Flux") generated_image = gr.Image(show_label=False, format="png", image_mode="RGBA", type="pil", height=300) gen3d_button = gr.Button("Create 3D visual with Trellis") video_output = gr.Video(label="Generated 3D Asset", autoplay=True, loop=True, height=300) message_box = gr.Textbox( label="Status Messages", interactive=False, placeholder="Status messages will appear here", ) # Accordion sections for advanced settings with gr.Accordion("Advanced Settings", open=False): with gr.Tab("Mistral"): # Mistral settings temperature = gr.Slider( label="Temperature", value=DEFAULT_TEMPERATURE, minimum=0.0, maximum=1.0, step=0.05, info="Higher values produce more diverse outputs", ) system_prompt = gr.Textbox( label="System Prompt", value=DEFAULT_SYSTEM_PROMPT, lines=10, info="Instructions for the Mistral model" ) with gr.Tab("Flux"): # Flux settings flux_seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=DEFAULT_SEED) flux_randomize_seed = gr.Checkbox(label="Randomize seed", value=DEFAULT_RANDOMIZE_SEED) with gr.Row(): width = gr.Slider(label="Width", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=DEFAULT_WIDTH) height = gr.Slider(label="Height", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=DEFAULT_HEIGHT) num_inference_steps = gr.Slider( label="Number of inference steps", minimum=1, maximum=50, step=1, value=DEFAULT_NUM_INFERENCE_STEPS, ) with gr.Tab("3D Generation Settings"): trellis_seed = gr.Slider(0, MAX_SEED, label="Seed", value=0, step=1) trellis_randomize_seed = gr.Checkbox(label="Randomize Seed", value=True) gr.Markdown("Stage 1: Sparse Structure Generation") with gr.Row(): ss_guidance_strength = gr.Slider(0.0, 10.0, label="Guidance Strength", value=7.5, step=0.1) ss_sampling_steps = gr.Slider(1, 50, label="Sampling Steps", value=12, step=1) gr.Markdown("Stage 2: Structured Latent Generation") with gr.Row(): slat_guidance_strength = gr.Slider(0.0, 10.0, label="Guidance Strength", value=3.0, step=0.1) slat_sampling_steps = gr.Slider(1, 50, label="Sampling Steps", value=12, step=1) gr.Examples( examples=examples, fn=process_example_pipeline, inputs=[prompt], outputs=[prompt], cache_examples=True, ) gr.on( triggers=[prompt_button.click, prompt.submit], fn=refine_prompt, inputs=[prompt, system_prompt], outputs=[refined_prompt, message_box] ) gr.on( triggers=[visual_button.click], fn=generate_image, inputs=[refined_prompt, flux_seed, flux_randomize_seed, width, height, num_inference_steps], outputs=[generated_image, message_box] ) gr.on( triggers=[gen3d_button.click], fn=image_to_3d, inputs=[generated_image, trellis_seed, ss_guidance_strength, ss_sampling_steps, slat_guidance_strength, slat_sampling_steps], outputs=[output_state, video_output], ) return demo if __name__ == "__main__": trellis = get_trellis_pipeline() trellis.cuda() demo = create_interface() demo.launch(debug=True)