Upload 2 files

app.py

@@ -8,6 +8,7 @@ import gradio as gr
 from pathlib import Path
 import tempfile
 import shutil
+from PIL import Image
 
 # Add the current directory to Python path
 sys.path.append(os.path.dirname(os.path.abspath(__file__)))
@@ -125,19 +126,83 @@ DEFAULT_CONFIG = {
     'log_light_power': 3.0
 }
 
-def process_garment(text_prompt, base_text_prompt, epochs, learning_rate, clip_weight, delta_clip_weight, progress=gr.Progress()):
+def process_garment(input_type, text_prompt, base_text_prompt, target_image, base_image, custom_mesh, epochs, learning_rate, clip_weight, delta_clip_weight, progress=gr.Progress()):
     """
     Main function to process garment generation
+    
+    Args:
+        input_type: Either "Text" or "Image" to determine the processing mode
+        text_prompt: Text description of target garment (for text mode)
+        base_text_prompt: Text description of base garment (for text mode)
+        target_image: Image of target garment style (for image mode)
+        base_image: Optional image of base garment (for image mode)
+        custom_mesh: Optional custom source mesh file (.obj)
+        epochs: Number of optimization epochs
+        learning_rate: Optimization learning rate
+        clip_weight: Weight for CLIP loss
+        delta_clip_weight: Weight for delta CLIP loss
+        progress: Gradio progress tracking object
     """
     try:
         # Create a temporary output directory
         with tempfile.TemporaryDirectory() as temp_dir:
             # Update configuration
             config = DEFAULT_CONFIG.copy()
+            
+            # Set up input parameters based on mode
+            if input_type == "Image" and target_image is not None:
+                # Image-based processing
+                progress(0.05, desc="Preparing image inputs...")
+                
+                # Save target image to temp directory
+                target_image_path = os.path.join(temp_dir, "target_image.jpg")
+                if isinstance(target_image, str):
+                    shutil.copy(target_image, target_image_path)
+                else:
+                    target_image.save(target_image_path)
+                
+                # Save base image if provided, or use default
+                if base_image is not None:
+                    base_image_path = os.path.join(temp_dir, "base_image.jpg")
+                    if isinstance(base_image, str):
+                        shutil.copy(base_image, base_image_path)
+                    else:
+                        base_image.save(base_image_path)
+                else:
+                    # If no base image, use a placeholder or default
+                    base_image_path = target_image_path
+                
+                # Configure for image-based processing
+                config.update({
+                    'image_prompt': target_image_path,
+                    'base_image_prompt': base_image_path,
+                    'fashion_image': True,
+                    'fashion_text': False
+                })
+            else:
+                # Text-based processing
+                if not text_prompt or len(text_prompt.strip()) == 0:
+                    return "Error: Text prompt is required for text-based generation."
+                    
+                if not base_text_prompt or len(base_text_prompt.strip()) == 0:
+                    base_text_prompt = "simple t-shirt"  # Default base prompt
+                    
+                config.update({
+                    'text_prompt': text_prompt,
+                    'base_text_prompt': base_text_prompt,
+                    'fashion_image': False,
+                    'fashion_text': True
+                })
+            
+            # Handle custom mesh if provided
+            if custom_mesh is not None:
+                custom_mesh_path = os.path.join(temp_dir, "custom_mesh.obj")
+                shutil.copy(custom_mesh, custom_mesh_path)
+                config['mesh'] = custom_mesh_path
+            
+            # Update optimization parameters
             config.update({
                 'output_path': temp_dir,
-                'text_prompt': text_prompt,
-                'base_text_prompt': base_text_prompt,
                 'epochs': int(epochs),
                 'lr': float(learning_rate),
                 'clip_weight': float(clip_weight),
@@ -160,14 +225,43 @@ def process_garment(text_prompt, base_text_prompt, epochs, learning_rate, clip_w
             
             progress(0.9, desc="Processing complete, preparing output...")
             
-            # Look for output files
-            output_files = []
+            # Look for output files, prioritize mesh files
+            obj_files = []
+            glb_files = []
+            image_files = []
+            
+            # First check for mesh files in mesh_final directory (priority)
+            mesh_final_dir = Path(temp_dir) / "mesh_final"
+            if mesh_final_dir.exists():
+                for file_path in mesh_final_dir.rglob("*"):
+                    if file_path.is_file():
+                        if file_path.suffix.lower() == '.obj':
+                            obj_files.append(str(file_path))
+                        elif file_path.suffix.lower() == '.glb':
+                            glb_files.append(str(file_path))
+            
+            # Check other mesh directories
+            for mesh_dir in Path(temp_dir).glob("mesh_*"):
+                if mesh_dir.is_dir() and mesh_dir.name != 'mesh_final':
+                    for file_path in mesh_dir.rglob("*"):
+                        if file_path.is_file():
+                            if file_path.suffix.lower() == '.obj':
+                                obj_files.append(str(file_path))
+                            elif file_path.suffix.lower() == '.glb':
+                                glb_files.append(str(file_path))
+            
+            # Collect image files for visualization
             for file_path in Path(temp_dir).rglob("*"):
-                if file_path.is_file() and file_path.suffix.lower() in ['.obj', '.png', '.jpg', '.jpeg', '.gif', '.mp4']:
-                    output_files.append(str(file_path))
+                if file_path.is_file() and file_path.suffix.lower() in ['.png', '.jpg', '.jpeg', '.gif', '.mp4']:
+                    image_files.append(str(file_path))
             
-            if output_files:
-                return output_files[0] if len(output_files) == 1 else output_files
+            # Prioritize output: GLB, OBJ, then images
+            if glb_files:
+                return glb_files[0]  # Return first GLB file (best for web viewing)
+            elif obj_files:
+                return obj_files[0]  # Return first OBJ file
+            elif image_files:
+                return image_files[0]  # Return an image if no mesh was found
             else:
                 return "Processing completed but no output files found."
                 
@@ -182,13 +276,15 @@ def create_interface():
         gr.Markdown("""
         # Garment3DGen: 3D Garment Stylization and Texture Generation
         
-        This tool allows you to stylize 3D garments using text prompts. Upload a 3D mesh and describe the desired style to generate a new 3D garment.
+        This tool allows you to stylize 3D garments using text prompts or images. Generate a new 3D garment mesh (.obj/.glb) 
+        that can be used for virtual try-on applications.
         
         ## How to use:
-        1. Enter a text prompt describing the target style (e.g., "leather jacket with studs")
-        2. Enter a base text prompt describing the input mesh (e.g., "simple t-shirt")
-        3. Adjust the parameters as needed
-        4. Click "Generate" to start the process
+        1. Choose your input method: Text Prompt or Image Upload
+        2. Enter parameters or upload images based on your selection
+        3. Optionally upload a custom mesh as your starting point
+        4. Adjust the optimization parameters as needed
+        5. Click "Generate" to start the process
         
         **Note:** Processing may take several minutes depending on the number of epochs.
         """)
@@ -197,76 +293,162 @@ def create_interface():
             with gr.Column(scale=1):
                 gr.Markdown("### Input Parameters")
                 
-                text_prompt = gr.Textbox(
-                    label="Target Text Prompt",
-                    placeholder="e.g., leather jacket with studs, denim jacket with patches",
-                    value="leather jacket with studs"
+                # Input type selector
+                input_type = gr.Radio(
+                    choices=["Text", "Image"],
+                    value="Text",
+                    label="Generation Method",
+                    info="Choose whether to generate from text descriptions or from images"
                 )
                 
-                base_text_prompt = gr.Textbox(
-                    label="Base Text Prompt",
-                    placeholder="e.g., simple t-shirt, basic long sleeve shirt",
-                    value="simple t-shirt"
-                )
+                # Text-based inputs
+                with gr.Group() as text_input_group:
+                    text_prompt = gr.Textbox(
+                        label="Target Text Prompt",
+                        placeholder="e.g., leather jacket with studs, denim jacket with patches",
+                        value="leather jacket with studs",
+                        info="Describe the garment style you want to create"
+                    )
+                    
+                    base_text_prompt = gr.Textbox(
+                        label="Base Text Prompt",
+                        placeholder="e.g., simple t-shirt, basic long sleeve shirt",
+                        value="simple t-shirt",
+                        info="Describe the starting garment style (default source mesh)"
+                    )
                 
-                epochs = gr.Slider(
-                    minimum=100,
-                    maximum=3000,
-                    value=1800,
-                    step=100,
-                    label="Number of Epochs",
-                    info="More epochs = better quality but longer processing time"
-                )
+                # Image-based inputs
+                with gr.Group(visible=False) as image_input_group:
+                    target_image = gr.Image(
+                        label="Target Garment Image",
+                        type="filepath",
+                        info="Upload an image of the garment style you want to create"
+                    )
+                    
+                    base_image = gr.Image(
+                        label="Base Garment Image (Optional)",
+                        type="filepath",
+                        info="Optionally upload an image of the starting garment style"
+                    )
                 
-                learning_rate = gr.Slider(
-                    minimum=0.0001,
-                    maximum=0.01,
-                    value=0.0025,
-                    step=0.0001,
-                    label="Learning Rate"
+                # Common inputs
+                custom_mesh = gr.File(
+                    label="Custom Source Mesh (Optional)",
+                    file_types=[".obj"],
+                    info="Optionally upload your own .obj file as the source mesh"
                 )
                 
-                clip_weight = gr.Slider(
-                    minimum=0.1,
-                    maximum=10.0,
-                    value=2.5,
-                    step=0.1,
-                    label="CLIP Weight"
-                )
+                # Advanced parameters section
+                with gr.Accordion("Advanced Parameters", open=False):
+                    epochs = gr.Slider(
+                        minimum=100,
+                        maximum=3000,
+                        value=1800,
+                        step=100,
+                        label="Number of Epochs",
+                        info="More epochs = better quality but longer processing time"
+                    )
+                    
+                    learning_rate = gr.Slider(
+                        minimum=0.0001,
+                        maximum=0.01,
+                        value=0.0025,
+                        step=0.0001,
+                        label="Learning Rate",
+                        info="Controls optimization speed and stability"
+                    )
+                    
+                    clip_weight = gr.Slider(
+                        minimum=0.1,
+                        maximum=10.0,
+                        value=2.5,
+                        step=0.1,
+                        label="CLIP Weight",
+                        info="Controls the influence of semantic features"
+                    )
+                    
+                    delta_clip_weight = gr.Slider(
+                        minimum=0.1,
+                        maximum=20.0,
+                        value=5.0,
+                        step=0.1,
+                        label="Delta CLIP Weight",
+                        info="Controls the strength of style transfer"
+                    )
                 
-                delta_clip_weight = gr.Slider(
-                    minimum=0.1,
-                    maximum=20.0,
-                    value=5.0,
-                    step=0.1,
-                    label="Delta CLIP Weight"
-                )
-                
-                generate_btn = gr.Button("Generate 3D Garment", variant="primary")
+                generate_btn = gr.Button("Generate 3D Garment", variant="primary", size="lg")
             
             with gr.Column(scale=1):
                 gr.Markdown("### Output")
-                output = gr.File(label="Generated 3D Garment")
-                status = gr.Textbox(label="Status", interactive=False)
+                output = gr.File(label="Generated 3D Garment (.obj/.glb)")
+                status = gr.Textbox(label="Status", value="Ready to generate", interactive=False)
+                
+                # Preview section
+                gr.Markdown("#### Mesh Preview")
+                gr.Markdown("""
+                The generated mesh can be:
+                - Imported into 3D software like Blender
+                - Used for virtual try-on applications
+                - Integrated with animation systems
+                - Used for digital fashion design
+                
+                Download the .obj file to use the mesh with your preferred 3D tools.
+                """)
+        
+        # Function to toggle between text and image inputs
+        def update_input_visibility(input_type):
+            if input_type == "Text":
+                return gr.Group.update(visible=True), gr.Group.update(visible=False)
+            else:  # Image
+                return gr.Group.update(visible=False), gr.Group.update(visible=True)
+        
+        # Set up the visibility toggle
+        input_type.change(
+            fn=update_input_visibility,
+            inputs=[input_type],
+            outputs=[text_input_group, image_input_group]
+        )
         
         # Connect the button to the processing function
         generate_btn.click(
             fn=process_garment,
-            inputs=[text_prompt, base_text_prompt, epochs, learning_rate, clip_weight, delta_clip_weight],
+            inputs=[
+                input_type,
+                text_prompt, 
+                base_text_prompt,
+                target_image,
+                base_image,
+                custom_mesh,
+                epochs, 
+                learning_rate, 
+                clip_weight, 
+                delta_clip_weight
+            ],
             outputs=[output]
         )
         
         gr.Markdown("""
         ## Tips for better results:
-        - Be specific in your text prompts
+        
+        ### Text Mode:
+        - Be specific in your text prompts (e.g., "red leather biker jacket with metal studs" is better than "jacket")
         - Use descriptive terms for materials, colors, and styles
         - The base text prompt should accurately describe your input mesh
+        
+        ### Image Mode:
+        - Use clear, well-lit images of garments
+        - Front-facing or 3/4 view images work best
+        - Images should show the complete garment
+        - Higher quality images produce better results
+        
+        ### General Tips:
         - Higher epoch counts generally produce better results but take longer
+        - Similar garment types work better (e.g., jacket→jacket works better than shirt→pants)
         - Experiment with different CLIP weights for different effects
         
         ## Technical Details:
-        This tool uses Neural Jacobian Fields and CLIP embeddings to deform and stylize 3D garment meshes.
-        The process involves optimizing the mesh geometry and texture to match the target text description.
+        This tool uses Neural Jacobian Fields and Fashion-CLIP embeddings to deform and stylize 3D garment meshes.
+        The process involves optimizing mesh geometry to match the target style while preserving garment structure.
         """)
     
     return interface
@@ -278,5 +460,8 @@ if __name__ == "__main__":
         server_name="0.0.0.0",
         server_port=7860,
         share=False,
-        debug=True
-    ) 
+        debug=False,
+        max_threads=16,  # Optimize thread usage for Hugging Face
+        show_error=True,  # Show detailed errors for better debugging
+        favicon_path="asset_visualization/armor.gif"  # Set a nice favicon
+    )