Update app.py

app.py

@@ -30,6 +30,8 @@ import trimesh
 import argparse
 import numpy as np
 import gradio as gr
+from gradio_client import Client
+from PIL import Image
 from step1x3d_geometry.models.pipelines.pipeline import Step1X3DGeometryPipeline
 from step1x3d_texture.pipelines.step1x_3d_texture_synthesis_pipeline import (
     Step1X3DTexturePipeline,
@@ -55,9 +57,46 @@ geometry_model = Step1X3DGeometryPipeline.from_pretrained(
 
 texture_model = Step1X3DTexturePipeline.from_pretrained("stepfun-ai/Step1X-3D", subfolder=args.texture_model)
 
+# Initialize text-to-image client
+t2i_client = Client("http://211.233.58.201:7971/")
+
+
+def generate_image_from_text(prompt, height, width, steps, scales, seed):
+    """Generate image from text using the external API"""
+    try:
+        result = t2i_client.predict(
+            height=height,
+            width=width,
+            steps=steps,
+            scales=scales,
+            prompt=prompt,
+            seed=seed if seed != -1 else None,
+            api_name="/process_and_save_image"
+        )
+        # Result contains a dict with 'path' key pointing to the generated image
+        if isinstance(result, dict) and 'path' in result:
+            return result['path']
+        elif isinstance(result, str):
+            return result
+        else:
+            raise Exception("Unexpected result format from text-to-image API")
+    except Exception as e:
+        print(f"Error generating image from text: {e}")
+        return None
+
+
+def get_random_seed():
+    """Get a random seed from the external API"""
+    try:
+        result = t2i_client.predict(api_name="/update_random_seed")
+        return result
+    except Exception as e:
+        print(f"Error getting random seed: {e}")
+        return -1
+
 
 @spaces.GPU(duration=240)
-def generate_func(
+def generate_3d_func(
     input_image_path, guidance_scale, inference_steps, max_facenum, symmetry, edge_type
 ):
     # geometry_model = geometry_model.to("cuda")
@@ -96,14 +135,46 @@ def generate_func(
     return geometry_save_path, textured_save_path
 
 
+def update_image_display(uploaded_image, generated_image):
+    """Update the displayed image based on which source has content"""
+    if generated_image is not None:
+        return generated_image
+    elif uploaded_image is not None:
+        return uploaded_image
+    else:
+        return None
+
+
 with gr.Blocks(title="Step1X-3D demo") as demo:
-    gr.Markdown("# Step1X-3D")
+    gr.Markdown("# Step1X-3D with Text-to-Image Generation")
+    
     with gr.Row():
         with gr.Column(scale=2):
-            input_image = gr.Image(label="Image", type="filepath")
-            guidance_scale = gr.Number(label="Guidance Scale", value="7.5")
+            gr.Markdown("## Image Input")
+            with gr.Tab("Upload Image"):
+                uploaded_image = gr.Image(label="Upload Image", type="filepath")
+            
+            with gr.Tab("Generate from Text"):
+                text_prompt = gr.Textbox(label="Image Description", placeholder="Enter your image description here...")
+                with gr.Row():
+                    t2i_height = gr.Slider(label="Height", minimum=512, maximum=2048, value=1024, step=64)
+                    t2i_width = gr.Slider(label="Width", minimum=512, maximum=2048, value=1024, step=64)
+                with gr.Row():
+                    t2i_steps = gr.Slider(label="Inference Steps", minimum=1, maximum=50, value=8, step=1)
+                    t2i_scales = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=10.0, value=3.5, step=0.5)
+                with gr.Row():
+                    t2i_seed = gr.Number(label="Seed (optional, -1 for random)", value=-1)
+                    random_seed_btn = gr.Button("Get Random Seed", scale=0)
+                generate_image_btn = gr.Button("Generate Image", variant="primary")
+            
+            # Display the current working image
+            current_image = gr.Image(label="Current Image (for 3D generation)", type="filepath", interactive=False)
+            generated_image_path = gr.State(value=None)
+            
+            gr.Markdown("## 3D Generation Settings")
+            guidance_scale = gr.Number(label="3D Guidance Scale", value="7.5")
             inference_steps = gr.Slider(
-                label="Inferece Steps", minimum=1, maximum=100, value=50
+                label="3D Inference Steps", minimum=1, maximum=100, value=50
             )
             max_facenum = gr.Number(label="Max Face Num", value="400000")
             symmetry = gr.Radio(
@@ -118,10 +189,12 @@ with gr.Blocks(title="Step1X-3D demo") as demo:
                 value="sharp",
                 type="value",
             )
-            btn = gr.Button("Start")
+            btn_3d = gr.Button("Generate 3D", variant="primary")
+            
         with gr.Column(scale=4):
             textured_preview = gr.Model3D(label="Textured", height=380)
             geometry_preview = gr.Model3D(label="Geometry", height=380)
+            
         with gr.Column(scale=1):
             gr.Examples(
                 examples=[
@@ -134,14 +207,58 @@ with gr.Blocks(title="Step1X-3D demo") as demo:
                     ["examples/images/061.png"],
                     ["examples/images/107.png"],
                 ],
-                inputs=[input_image],
+                inputs=[uploaded_image],
                 cache_examples=False,
+                label="Example Images"
             )
 
-    btn.click(
-        generate_func,
+    # Event handlers
+    def on_generate_image(prompt, height, width, steps, scales, seed):
+        if not prompt:
+            gr.Warning("Please enter a text prompt")
+            return None, None
+        
+        generated_path = generate_image_from_text(prompt, height, width, steps, scales, seed)
+        if generated_path:
+            return generated_path, generated_path
+        else:
+            gr.Warning("Failed to generate image from text")
+            return None, None
+    
+    def on_upload_image(image_path):
+        return image_path
+    
+    def get_current_image(uploaded, generated):
+        if generated is not None:
+            return generated
+        elif uploaded is not None:
+            return uploaded
+        else:
+            return None
+    
+    # Connect event handlers
+    generate_image_btn.click(
+        on_generate_image,
+        inputs=[text_prompt, t2i_height, t2i_width, t2i_steps, t2i_scales, t2i_seed],
+        outputs=[generated_image_path, current_image]
+    )
+    
+    random_seed_btn.click(
+        get_random_seed,
+        inputs=[],
+        outputs=[t2i_seed]
+    )
+    
+    uploaded_image.change(
+        on_upload_image,
+        inputs=[uploaded_image],
+        outputs=[current_image]
+    )
+    
+    btn_3d.click(
+        lambda img, gs, is_, mf, sym, et: generate_3d_func(img, gs, is_, mf, sym, et) if img else (None, None),
         inputs=[
-            input_image,
+            current_image,
             guidance_scale,
             inference_steps,
             max_facenum,
@@ -151,4 +268,4 @@ with gr.Blocks(title="Step1X-3D demo") as demo:
         outputs=[geometry_preview, textured_preview],
     )
 
-demo.launch(ssr_mode=False)
+demo.launch(ssr_mode=False)