Create app.py

app.py

@@ -0,0 +1,419 @@
+import gradio as gr
+import os
+import zipfile
+import time
+import uuid # For unique filenames
+
+# --- LLM/Model Setup ---
+from transformers import pipeline as transformers_pipeline # For local list generation
+from huggingface_hub import InferenceClient # For prompt refinement via API
+from diffusers import StableDiffusionPipeline, EulerDiscreteScheduler # For image generation
+from gradio_client import Client as GradioClient, handle_file # For 3D generation
+
+# --- Configuration ---
+# Consider making these configurable in the UI later
+LIST_GENERATION_MODEL = "google/flan-t5-base" # Or another suitable small model
+PROMPT_REFINEMENT_MODEL_API = "mistralai/Mixtral-8x7B-Instruct-v0.1" # Or another instruct model via Inference API
+IMAGE_GENERATION_MODEL = "stabilityai/stable-diffusion-xl-base-1.0" # Or "runwayml/stable-diffusion-v1-5"
+HUNYUAN_SPACE_ID = "tencent/Hunyuan3D-2"
+OUTPUT_DIR = "outputs"
+MODELS_SUBDIR = "3d_models"
+IMAGES_SUBDIR = "image_previews"
+ZIP_FILENAME = "3d_collection.zip"
+
+# --- Initialize Clients/Pipelines (can be slow, consider loading on demand if needed) ---
+
+# Use HF Token from Space secrets if available/needed for Inference API
+HF_TOKEN = os.environ.get("HUGGINGFACE_TOKEN", None)
+
+# Basic List Generator (local)
+try:
+    list_generator = transformers_pipeline("text2text-generation", model=LIST_GENERATION_MODEL)
+except Exception as e:
+    print(f"Warning: Could not load local list generator {LIST_GENERATION_MODEL}: {e}")
+    list_generator = None
+
+# Prompt Refiner (API)
+try:
+    if not HF_TOKEN:
+        print("Warning: HUGGINGFACE_TOKEN not set. Inference API calls might be rate-limited or fail.")
+    prompt_refiner = InferenceClient(model=PROMPT_REFINEMENT_MODEL_API, token=HF_TOKEN)
+except Exception as e:
+    print(f"Warning: Could not initialize InferenceClient for {PROMPT_REFINEMENT_MODEL_API}: {e}")
+    prompt_refiner = None
+
+# Image Generator (Local with Diffusers - requires GPU on Space for reasonable speed)
+# Or consider an Image Gen API service if running on CPU hardware
+try:
+    # Using XL as an example - adjust based on available hardware
+    image_pipeline = StableDiffusionPipeline.from_pretrained(IMAGE_GENERATION_MODEL, torch_dtype=torch.float16, use_safetensors=True)
+    # Move to GPU if available (check Space hardware)
+    # image_pipeline.to("cuda") # Uncomment if GPU is available
+    image_pipeline.scheduler = EulerDiscreteScheduler.from_config(image_pipeline.scheduler.config)
+except Exception as e:
+    print(f"Warning: Could not load diffusers pipeline {IMAGE_GENERATION_MODEL}. Image generation might fail: {e}")
+    image_pipeline = None
+
+
+# 3D Generator Client
+try:
+    hunyuan_client = GradioClient(HUNYUAN_SPACE_ID)
+except Exception as e:
+    print(f"Error initializing GradioClient for {HUNYUAN_SPACE_ID}: {e}")
+    hunyuan_client = None
+
+# --- Helper Functions ---
+
+def generate_list_local(theme, count):
+    if not list_generator:
+        return ["Error: List generator model not loaded."]
+    prompt = f"Generate a comma-separated list of {count} distinct types of {theme}."
+    try:
+        result = list_generator(prompt, max_length=200)[0]['generated_text']
+        items = [item.strip() for item in result.split(',') if item.strip()]
+        return items[:count] # Ensure we don't exceed the requested count
+    except Exception as e:
+        print(f"Error generating list: {e}")
+        return [f"Error: {e}"]
+
+def refine_prompt_api(item_name):
+    if not prompt_refiner:
+        return f"A 3D model of a {item_name}" # Fallback basic prompt
+    prompt = f"Create a detailed, descriptive prompt for generating a highly realistic image of a single '{item_name}'. Focus on visual details suitable for a text-to-image AI. Only output the prompt itself."
+    try:
+        refined = prompt_refiner.text_generation(prompt, max_new_tokens=100)
+        # Clean up potential API artifacts if necessary
+        refined = refined.strip().strip('"')
+        return refined
+    except Exception as e:
+        print(f"Error refining prompt for '{item_name}': {e}")
+        # Fallback to a simpler prompt for 3D generation if refinement fails
+        return f"A high quality 3D model of a {item_name}"
+
+def generate_image_local(refined_prompt, output_path):
+    if not image_pipeline:
+        print("Image generation pipeline not available.")
+        # Create a placeholder image or return None
+        # Example: from PIL import Image; img = Image.new('RGB', (60, 30), color = 'red'); img.save(output_path); return output_path
+        return None
+    try:
+        # Adjust inference steps/guidance as needed
+        image = image_pipeline(refined_prompt, num_inference_steps=25, guidance_scale=7.5).images[0]
+        image.save(output_path)
+        return output_path
+    except Exception as e:
+        print(f"Error generating image for prompt '{refined_prompt}': {e}")
+        return None
+
+def generate_3d_model_hunyuan(refined_prompt_for_3d, output_dir, item_name_safe):
+    if not hunyuan_client:
+        print("Hunyuan 3D client not available.")
+        return None, "Client not initialized"
+
+    print(f"Requesting 3D model for: {refined_prompt_for_3d}")
+    # Use defaults for most parameters initially
+    try:
+        result_tuple = hunyuan_client.predict(
+                caption=refined_prompt_for_3d,
+                # Leave image and mv_image inputs as None for text-to-3D
+                image=None,
+                mv_image_front=None,
+                mv_image_back=None,
+                mv_image_left=None,
+                mv_image_right=None,
+                # Default values from API docs (can be overridden)
+                steps=30,
+                guidance_scale=5,
+                seed=1234, # Or use randomize_seed=True
+                octree_resolution=256,
+                check_box_rembg=True,
+                num_chunks=8000,
+                randomize_seed=True,
+                api_name="/generation_all" # Crucial!
+        )
+
+        # --- VERIFICATION NEEDED ---
+        # Check the actual return tuple structure. Assuming file path is first or second.
+        # Let's try the first element (index 0). If it's None or not a path, try index 1.
+        raw_filepath = None
+        if result_tuple and len(result_tuple) > 0 and isinstance(result_tuple[0], str):
+             raw_filepath = result_tuple[0]
+        elif result_tuple and len(result_tuple) > 1 and isinstance(result_tuple[1], str):
+             print("Using second element from result tuple for filepath.")
+             raw_filepath = result_tuple[1]
+        # --- END VERIFICATION NEEDED ---
+
+        if raw_filepath:
+            print(f"Job completed. Raw result path: {raw_filepath}")
+            os.makedirs(output_dir, exist_ok=True)
+
+            # Download the file using handle_file which manages temp paths etc.
+            # handle_file saves with a potentially random name in download_dir
+            downloaded_temp_path = handle_file(raw_filepath, download_dir=output_dir)
+
+            if downloaded_temp_path and os.path.exists(downloaded_temp_path):
+                # Rename it to something meaningful
+                file_ext = os.path.splitext(downloaded_temp_path)[1] # Get extension (.glb, .obj?)
+                if not file_ext: file_ext = ".glb" # Assume glb if unknown
+                final_path = os.path.join(output_dir, f"{item_name_safe}{file_ext}")
+                os.rename(downloaded_temp_path, final_path)
+                print(f"Model saved to: {final_path}")
+                return final_path, "Success"
+            else:
+                 error_msg = f"handle_file failed to download or returned invalid path: {downloaded_temp_path}"
+                 print(error_msg)
+                 return None, error_msg
+        else:
+            error_msg = f"Job for '{refined_prompt_for_3d}' did not return a valid filepath in expected tuple elements."
+            print(error_msg)
+            # You might want to inspect the full result_tuple here for debugging
+            print(f"Full result tuple: {result_tuple}")
+            return None, error_msg
+
+    except Exception as e:
+        error_msg = f"Error calling Hunyuan3D API for '{refined_prompt_for_3d}': {e}"
+        print(error_msg)
+        return None, str(e)
+
+def create_zip(files_to_zip, zip_filepath):
+    with zipfile.ZipFile(zip_filepath, 'w') as zf:
+        for file_path in files_to_zip:
+            if file_path and os.path.exists(file_path):
+                zf.write(file_path, os.path.basename(file_path))
+    return zip_filepath
+
+
+# --- Gradio Interface & Logic ---
+
+with gr.Blocks() as demo:
+    gr.Markdown("# 3D Asset Collection Generator")
+    gr.Markdown("Generate a list based on a theme, refine prompts, preview images, and generate selected 3D models using Hunyuan3D-2.")
+    if not HF_TOKEN:
+        gr.Warning("Hugging Face Token not found. Prompt refinement quality/rate limits may be affected. Consider adding HUGGINGFACE_TOKEN to Space secrets.")
+    if not image_pipeline:
+        gr.Warning("Local Image Generation model failed to load. Image previews will be skipped. Check Space hardware/logs.")
+    if not hunyuan_client:
+        gr.Error("Failed to connect to the Hunyuan3D-2 Space. 3D generation will not work.")
+
+
+    # State to hold intermediate results
+    # Using gr.State is good for simple values, for complex lists/dicts might need alternatives or careful handling
+    list_items_state = gr.State([])
+    refined_prompts_state = gr.State({}) # Dict: {item_name: refined_prompt}
+    image_paths_state = gr.State({}) # Dict: {item_name: image_path}
+    selected_items_state = gr.State([]) # List of item_names selected by user
+    generated_3d_files_state = gr.State([]) # List of paths to successfully generated models
+
+
+    with gr.Row():
+        theme_input = gr.Textbox(label="Theme", placeholder="e.g., reptiles, kitchen appliances, medieval weapons")
+        count_input = gr.Number(label="Number of Items", value=5, minimum=1, step=1)
+
+    generate_list_button = gr.Button("1. Generate List & Refine Prompts")
+    list_output_display = gr.Markdown("List will appear here...") # Or use gr.DataFrame
+
+    generate_images_button = gr.Button("2. Generate Image Previews", visible=False) # Hidden initially
+    # Use Gallery for display, Dataset for selection tracking
+    image_gallery = gr.Gallery(label="Image Previews", visible=False, elem_id="image_gallery")
+    # Dataset to hold data for selection (item_name, image_path, refined_prompt)
+    selection_data = gr.Dataset(components=[gr.Textbox(visible=False), gr.Textbox(visible=False), gr.Textbox(visible=False)], # item, img_path, prompt
+                                 headers=["Item Name", "Image", "Prompt"],
+                                 label="Select Items for 3D Generation",
+                                 visible=False)
+
+
+    generate_3d_button = gr.Button("3. Generate 3D Models for Selected Items", visible=False) # Hidden initially
+    status_output = gr.Markdown("") # For progress updates
+    final_zip_output = gr.File(label="Download 3D Model Collection (ZIP)", visible=False)
+
+
+    # --- Event Logic ---
+
+    def run_list_and_refine(theme, count):
+        if not theme:
+            return {list_output_display: "Please enter a theme.", generate_images_button: gr.Button(visible=False)}
+
+        # Ensure output dirs exist
+        os.makedirs(os.path.join(OUTPUT_DIR, IMAGES_SUBDIR), exist_ok=True)
+        os.makedirs(os.path.join(OUTPUT_DIR, MODELS_SUBDIR), exist_ok=True)
+
+
+        gr.Info("Generating list...")
+        items = generate_list_local(theme, int(count))
+        if not items or "Error:" in items[0]:
+            return {list_output_display: f"Failed to generate list: {items[0] if items else 'Unknown error'}",
+                    generate_images_button: gr.Button(visible=False)}
+
+        list_items_state.value = items # Save items to state
+
+        gr.Info("Refining prompts via API...")
+        refined_prompts = {}
+        output_md = "### Generated List & Refined Prompts:\n\n"
+        for item in items:
+            refined = refine_prompt_api(item)
+            refined_prompts[item] = refined
+            output_md += f"*   **{item}:** {refined}\n"
+
+        refined_prompts_state.value = refined_prompts # Save refined prompts
+
+        # Enable next step
+        return {
+            list_output_display: output_md,
+            generate_images_button: gr.Button(visible=True) # Show image gen button
+        }
+
+    generate_list_button.click(
+        fn=run_list_and_refine,
+        inputs=[theme_input, count_input],
+        outputs=[list_output_display, generate_images_button, list_items_state, refined_prompts_state] # Update state too
+    )
+
+
+    def run_image_generation(items, refined_prompts_dict):
+        if not image_pipeline:
+             # Skip image generation if pipeline not loaded
+             gr.Warning("Image pipeline not loaded. Skipping image previews.")
+             # Prepare data for selection without images
+             selection_samples = [[item, "N/A", refined_prompts_dict.get(item, "")] for item in items]
+             image_paths_state.value = {} # Clear image paths
+             return {
+                 image_gallery: gr.Gallery(visible=False),
+                 selection_data: gr.Dataset(samples=selection_samples, visible=True),
+                 generate_3d_button: gr.Button(visible=True) # Allow proceeding without previews
+             }
+
+        gr.Info("Generating image previews... (this may take a while)")
+        image_paths = {}
+        gallery_images = []
+        selection_samples = [] # For the Dataset component
+
+        img_dir = os.path.join(OUTPUT_DIR, IMAGES_SUBDIR)
+
+        for item in items:
+            refined_prompt = refined_prompts_dict.get(item, f"Image of {item}") # Get refined prompt
+            safe_item_name = "".join(c if c.isalnum() else "_" for c in item)
+            img_filename = f"{safe_item_name}_{uuid.uuid4()}.png"
+            img_path = os.path.join(img_dir, img_filename)
+
+            generated_path = generate_image_local(refined_prompt, img_path)
+
+            if generated_path:
+                image_paths[item] = generated_path
+                gallery_images.append(generated_path)
+                selection_samples.append([item, generated_path, refined_prompt])
+            else:
+                # Handle image generation failure - maybe add placeholder info
+                selection_samples.append([item, "Failed", refined_prompt])
+                # Optionally add a placeholder to gallery_images too
+
+        image_paths_state.value = image_paths # Save image paths
+
+        # Show gallery and selection dataset
+        return {
+            image_gallery: gr.Gallery(value=gallery_images, visible=True),
+            selection_data: gr.Dataset(samples=selection_samples, visible=True),
+            generate_3d_button: gr.Button(visible=True) # Show 3D gen button
+        }
+
+    generate_images_button.click(
+        fn=run_image_generation,
+        inputs=[list_items_state, refined_prompts_state],
+        outputs=[image_gallery, selection_data, generate_3d_button, image_paths_state] # Update state
+    )
+
+    # Handler for when user makes selections in the Dataset
+    # Note: Gradio's Dataset selection handling might require specific event listeners
+    # or potentially using gr.CheckboxGroup or similar if Dataset selection is tricky.
+    # For simplicity here, we assume we can get the selected indices/items.
+    # A common pattern is to add a hidden Textbox updated by JS on selection,
+    # or use the Dataset's 'select' event if available and robust.
+    # Let's simulate getting selected *items* (requires correct component setup).
+    # This part might need refinement based on Gradio version/behavior.
+
+    # We'll trigger 3D generation directly from the button click for now,
+    # assuming the selection_data component holds the necessary info and selection state.
+
+    def run_3d_generation(selection_evt: gr.SelectData, all_items_data):
+        if not hunyuan_client:
+            return {status_output: "Hunyuan3D client not initialized. Cannot generate.", final_zip_output: gr.File(visible=False)}
+
+        selected_indices = selection_evt.index if selection_evt else []
+        if not selected_indices:
+            return {status_output: "Please select items from the table above before generating 3D models.", final_zip_output: gr.File(visible=False)}
+
+        # Extract selected items based on indices from the *current* data in the dataset
+        selected_items_info = [all_items_data[i] for i in selected_indices] # Each item is [name, img_path, prompt]
+
+        generated_files = []
+        status_messages = ["### 3D Generation Status:\n"]
+
+        model_dir = os.path.join(OUTPUT_DIR, MODELS_SUBDIR)
+
+        total_selected = len(selected_items_info)
+        for i, (item_name, _, refined_prompt) in enumerate(selected_items_info):
+            current_status = f"({i+1}/{total_selected}) Generating model for: **{item_name}**..."
+            print(current_status)
+            status_messages.append(f"*   {current_status}")
+            # Update UI status progressively
+            yield {status_output: "\n".join(status_messages), final_zip_output: gr.File(visible=False)}
+
+            # Adapt prompt slightly for 3D if desired, or use the image prompt directly
+            prompt_for_3d = refined_prompt # Or customize: f"A high quality 3D model of {item_name}, {refined_prompt}"
+
+            item_name_safe = "".join(c if c.isalnum() else "_" for c in item_name)
+
+            # --- Retry Logic Placeholder ---
+            max_retries = 1 # Example: allow 1 retry
+            attempts = 0
+            model_path = None
+            last_error = "Unknown error"
+
+            while attempts <= max_retries:
+                attempts += 1
+                if attempts > 1:
+                     status_messages.append(f"    *   Retrying ({attempts-1}/{max_retries})...")
+                     yield {status_output: "\n".join(status_messages)}
+                     time.sleep(2) # Brief pause before retry
+
+                model_path, msg = generate_3d_model_hunyuan(prompt_for_3d, model_dir, item_name_safe)
+                last_error = msg
+                if model_path:
+                    generated_files.append(model_path)
+                    status_messages.append(f"    *   Success! Model saved.")
+                    break # Exit retry loop on success
+                else:
+                    status_messages.append(f"    *   Attempt {attempts} failed: {msg}")
+
+            if not model_path:
+                 status_messages.append(f"    *   **Failed** after {attempts} attempt(s). Last error: {last_error}")
+            # --- End Retry Logic ---
+
+            # Update UI status after each item
+            yield {status_output: "\n".join(status_messages)}
+
+
+        if generated_files:
+            status_messages.append("\nCreating ZIP archive...")
+            yield {status_output: "\n".join(status_messages)}
+            zip_path = os.path.join(OUTPUT_DIR, ZIP_FILENAME)
+            final_zip = create_zip(generated_files, zip_path)
+            status_messages.append(f"\n**Collection ready!** Download '{ZIP_FILENAME}' below.")
+            generated_3d_files_state.value = generated_files # Store final paths
+            return {status_output: "\n".join(status_messages), final_zip_output: gr.File(value=final_zip, visible=True)}
+        else:
+            status_messages.append("\nNo 3D models were successfully generated.")
+            return {status_output: "\n".join(status_messages), final_zip_output: gr.File(visible=False)}
+
+
+    # Link the button click to the generator function
+    # The 'select' event on Dataset provides selection info (gr.SelectData)
+    # We pass both the selection event data and the full dataset content
+    generate_3d_button.click(
+        fn=run_3d_generation,
+        inputs=[selection_data, selection_data], # Pass dataset twice: once for select event, once for full data access
+        outputs=[status_output, final_zip_output, generated_3d_files_state] # Update state
+    )
+
+
+# Launch the Gradio app
+demo.queue().launch(debug=True) # Enable queue for longer processes, debug for detailed errors