Fix 3D generation

README.md

@@ -26,9 +26,9 @@ This Hugging Face Space provides a complete workflow to generate 3D models from
 ## Features
 
 - 🎨 **Text to Image**: Generate images using Stable Diffusion XL base model
+- 📤 **Flexible Input**: Use generated images or upload your own images for 3D generation
 - ✂️ **Automatic Background Removal**: Background removal is handled automatically by Stable Fast 3D
-- 🎮 **3D Generation**: Create textured 3D mesh models from images
-- 🔄 **Step-by-step Workflow**: Review and confirm at each step
+- 🎮 **3D Generation**: Create textured 3D mesh models from images with live preview
 - ⚙️ **Customizable**: Adjust remeshing options, vertex count, and texture resolution
 
 ## How to Use
@@ -40,14 +40,15 @@ This Hugging Face Space provides a complete workflow to generate 3D models from
    - Click "Generate Image" and wait for the result
 
 2. **Step 2 - 3D Generation**:
-   - Review the generated image
+   - The generated image from Step 1 is automatically filled in the input field
+   - **OR** upload your own image to skip Step 1 entirely
    - Adjust 3D generation settings:
      - **Remeshing Option**: Choose "none", "triangle", or "quad" remeshing
      - **Target Vertex Count**: Set to -1 for automatic, or specify a target count
      - **Texture Size**: Choose texture resolution (512-2048)
-   - Click "Continue to 3D Generation" to create the 3D model
+   - Click "Generate 3D Model" to create the 3D model
    - Background removal is handled automatically by Stable Fast 3D
-   - Download your GLB file
+   - View the 3D model preview and download your GLB file
 
 ## Tips
 

app.py

@@ -265,23 +265,26 @@ def generate_3d_from_image(
 def step1_generate_image(prompt, negative_prompt, num_steps):
     """Step 1: Generate image from text."""
     if not prompt:
-        return None, gr.update(visible=False), "Please enter a prompt"
+        return None, None, "Please enter a prompt"
 
     try:
         image = generate_text_to_image(prompt, negative_prompt, num_steps)
         return (
             image,
-            gr.update(visible=True, value="Continue to 3D Generation"),
-            "Image generated successfully! Review and continue to generate 3D model.",
+            image,  # Auto-fill Step 2 image input
+            "Image generated successfully! The image has been automatically filled in Step 2. You can also upload your own image in Step 2.",
         )
     except Exception as e:
-        return None, gr.update(visible=False), f"Error generating image: {str(e)}"
+        return None, None, f"Error generating image: {str(e)}"
 
 
 def step2_generate_3d(image, remesh_option, vertex_count, texture_size):
     """Step 2: Generate 3D model from image (with built-in background removal)."""
     if image is None:
-        return gr.update(value=None, visible=False), "Please generate an image first"
+        return (
+            gr.update(value=None, visible=False),
+            "Please provide an image (generate one in Step 1 or upload your own)",
+        )
 
     try:
         glb_file = generate_3d_from_image(
@@ -290,7 +293,7 @@ def step2_generate_3d(image, remesh_option, vertex_count, texture_size):
 
         return (
             gr.update(value=glb_file, visible=True),
-            "3D model generated successfully! Background removal was handled automatically. You can download the model below.",
+            "3D model generated successfully! Background removal was handled automatically. You can view and download the model below.",
         )
     except Exception as e:
         return (
@@ -340,20 +343,22 @@ with gr.Blocks(title="Text to Image to 3D") as demo:
             )
             generate_btn = gr.Button("Generate Image", variant="primary")
             step1_status = gr.Textbox(label="Status", interactive=False)
-            step1_continue_btn = gr.Button(
-                "Continue to 3D Generation",
-                visible=False,
-                variant="secondary",
-            )
 
         with gr.Column(scale=1):
             step1_image = gr.Image(label="Generated Image", type="pil")
 
     # Step 2: 3D Generation
     gr.Markdown("## Step 2: 3D Generation")
-    gr.Markdown("*Background removal is handled automatically*")
+    gr.Markdown(
+        "*Background removal is handled automatically. You can use the image from Step 1 or upload your own image.*"
+    )
     with gr.Row():
         with gr.Column(scale=1):
+            step2_image_input = gr.Image(
+                label="Input Image",
+                type="pil",
+                sources=["upload", "clipboard"],
+            )
             remesh_option = gr.Radio(
                 choices=["none", "triangle", "quad"],
                 label="Remeshing Option",
@@ -378,8 +383,8 @@ with gr.Blocks(title="Text to Image to 3D") as demo:
 
         with gr.Column(scale=1):
             step2_output = LitModel3D(
-                label="3D Model",
-                visible=False,
+                label="3D Model Preview",
+                visible=True,
                 clear_color=[0.0, 0.0, 0.0, 0.0],
             )
 
@@ -387,12 +392,12 @@ with gr.Blocks(title="Text to Image to 3D") as demo:
     generate_btn.click(
         fn=step1_generate_image,
         inputs=[prompt, negative_prompt, num_steps],
-        outputs=[step1_image, step1_continue_btn, step1_status],
+        outputs=[step1_image, step2_image_input, step1_status],
     )
 
-    step1_continue_btn.click(
+    step2_generate_btn.click(
         fn=step2_generate_3d,
-        inputs=[step1_image, remesh_option, vertex_count, texture_size],
+        inputs=[step2_image_input, remesh_option, vertex_count, texture_size],
         outputs=[step2_output, step2_status],
     )