Update app.py

app.py

@@ -5,7 +5,7 @@ import imageio
 import numpy as np
 import torch
 import rembg
-from PIL import Image
+from PIL import Image, ImageDraw, ImageFont
 from torchvision.transforms import v2
 from pytorch_lightning import seed_everything
 from omegaconf import OmegaConf
@@ -192,35 +192,6 @@ def make3d(images):
         # get triplane
         planes = model.forward_planes(images, input_cameras)
 
-        # # get video
-        # chunk_size = 20 if IS_FLEXICUBES else 1
-        # render_size = 384
-        
-        # frames = []
-        # for i in tqdm(range(0, render_cameras.shape[1], chunk_size)):
-        #     if IS_FLEXICUBES:
-        #         frame = model.forward_geometry(
-        #             planes,
-        #             render_cameras[:, i:i+chunk_size],
-        #             render_size=render_size,
-        #         )['img']
-        #     else:
-        #         frame = model.synthesizer(
-        #             planes,
-        #             cameras=render_cameras[:, i:i+chunk_size],
-        #             render_size=render_size,
-        #         )['images_rgb']
-        #     frames.append(frame)
-        # frames = torch.cat(frames, dim=1)
-
-        # images_to_video(
-        #     frames[0],
-        #     video_fpath,
-        #     fps=30,
-        # )
-
-        # print(f"Video saved to {video_fpath}")
-
         # get mesh
         mesh_out = model.extract_mesh(
             planes,
@@ -239,6 +210,37 @@ def make3d(images):
     return mesh_fpath, mesh_glb_fpath
 
 
+# New function to generate 2D pixel art sprites
+def generate_pixel_art(prompt, remove_background=True, sample_steps=75, seed=42):
+    """Generate a pixel art sprite based on the prompt"""
+    seed_everything(seed)
+    
+    # Create a simple image with text as starting point
+    text_img = Image.new('RGB', (512, 512), color=(255, 255, 255))
+    draw = ImageDraw.Draw(text_img)
+    # Try to load a font, use default if not available
+    try:
+        font = ImageFont.truetype("Arial", 20)
+    except:
+        font = ImageFont.load_default()
+    
+    # Add prompt as text
+    pixel_prompt = f"Pixel art: {prompt}"
+    draw.text((10, 10), pixel_prompt, fill=(0, 0, 0), font=font)
+    
+    # Process through the pipeline
+    processed_img = preprocess(text_img, remove_background)
+    
+    # Generate the pixel art
+    result_img, _ = generate_mvs(processed_img, sample_steps, seed)
+    
+    # Save to a temporary file
+    sprite_path = tempfile.NamedTemporaryFile(suffix=".png", delete=False).name
+    result_img.save(sprite_path)
+    
+    return sprite_path
+
+
 _HEADER_ = '''
 <h2><b>Official 🤗 Gradio Demo</b></h2><h2><a href='https://github.com/TencentARC/InstantMesh' target='_blank'><b>InstantMesh: Efficient 3D Mesh Generation from a Single Image with Sparse-view Large Reconstruction Models</b></a></h2>
 
@@ -277,98 +279,96 @@ If you have any questions, feel free to open a discussion or contact us at <b>bl
 
 
 with gr.Blocks() as demo:
-    gr.Markdown(_HEADER_)
-    with gr.Row(variant="panel"):
-        with gr.Column():
-            with gr.Row():
-                input_image = gr.Image(
-                    label="Input Image",
-                    image_mode="RGBA",
-                    sources="upload",
-                    #width=256,
-                    #height=256,
-                    type="pil",
-                    elem_id="content_image",
-                )
-                processed_image = gr.Image(
-                    label="Processed Image", 
-                    image_mode="RGBA", 
-                    #width=256,
-                    #height=256,
-                    type="pil", 
-                    interactive=False
-                )
-            with gr.Row():
-                with gr.Group():
-                    do_remove_background = gr.Checkbox(
-                        label="Remove Background", value=True
+    gr.Markdown("# InstantMesh and Pixel Art Generator")
+    
+    with gr.Tab("3D Model Generation"):
+        with gr.Row(variant="panel"):
+            with gr.Column():
+                with gr.Row():
+                    input_image = gr.Image(
+                        label="Input Image",
+                        image_mode="RGBA",
+                        sources="upload",
+                        type="pil",
+                        elem_id="content_image",
                     )
-                    sample_seed = gr.Number(value=42, label="Seed Value", precision=0)
-
-                    sample_steps = gr.Slider(
+                    processed_image = gr.Image(
+                        label="Processed Image", 
+                        image_mode="RGBA", 
+                        type="pil", 
+                        interactive=False
+                    )
+                with gr.Row():
+                    with gr.Group():
+                        do_remove_background = gr.Checkbox(
+                            label="Remove Background", value=True
+                        )
+                        sample_seed = gr.Number(value=42, label="Seed Value", precision=0)
+                        sample_steps = gr.Slider(
+                            label="Sample Steps",
+                            minimum=30,
+                            maximum=75,
+                            value=75,
+                            step=5
+                        )
+
+                with gr.Row():
+                    submit_3d = gr.Button("Generate 3D Model", elem_id="generate", variant="primary")
+
+            with gr.Column():
+                with gr.Row():
+                    with gr.Column():
+                        mv_show_images = gr.Image(
+                            label="Generated Multi-views",
+                            type="pil",
+                            width=379,
+                            interactive=False
+                        )
+
+                with gr.Row():
+                    with gr.Tab("OBJ"):
+                        output_model_obj = gr.Model3D(
+                            label="Output Model (OBJ Format)",
+                            interactive=False,
+                        )
+                    with gr.Tab("GLB"):
+                        output_model_glb = gr.Model3D(
+                            label="Output Model (GLB Format)",
+                            interactive=False,
+                        )
+    
+    with gr.Tab("Pixel Art Generation"):
+        with gr.Row(variant="panel"):
+            with gr.Column():
+                pixel_prompt = gr.Textbox(
+                    label="Describe your pixel art sprite",
+                    placeholder="green cactus snake character, side view, game sprite",
+                    lines=3
+                )
+                with gr.Row():
+                    pixel_bg_remove = gr.Checkbox(label="Remove Background", value=True)
+                    pixel_seed = gr.Number(value=42, label="Seed Value", precision=0)
+                    pixel_steps = gr.Slider(
                         label="Sample Steps",
                         minimum=30,
                         maximum=75,
                         value=75,
                         step=5
                     )
-
-            with gr.Row():
-                submit = gr.Button("Generate", elem_id="generate", variant="primary")
-
-            with gr.Row(variant="panel"):
-                gr.Examples(
-                    examples=[
-                        os.path.join("examples", img_name) for img_name in sorted(os.listdir("examples"))
-                    ],
-                    inputs=[input_image],
-                    label="Examples",
-                    cache_examples=False,
-                    examples_per_page=16
+                submit_pixel = gr.Button("Generate Pixel Art", variant="primary")
+            
+            with gr.Column():
+                pixel_output = gr.Image(
+                    label="Generated Pixel Art Sprite",
+                    type="pil",
+                    interactive=False
                 )
-
-        with gr.Column():
-
-            with gr.Row():
-
-                with gr.Column():
-                    mv_show_images = gr.Image(
-                        label="Generated Multi-views",
-                        type="pil",
-                        width=379,
-                        interactive=False
-                    )
-
-                # with gr.Column():
-                #     output_video = gr.Video(
-                #         label="video", format="mp4",
-                #         width=379,
-                #         autoplay=True,
-                #         interactive=False
-                #     )
-
-            with gr.Row():
-                with gr.Tab("OBJ"):
-                    output_model_obj = gr.Model3D(
-                        label="Output Model (OBJ Format)",
-                        interactive=False,
-                    )
-                    gr.Markdown("Note: Downloaded .obj model will be flipped. Export .glb instead or manually flip it before usage.")
-                with gr.Tab("GLB"):
-                    output_model_glb = gr.Model3D(
-                        label="Output Model (GLB Format)",
-                        interactive=False,
-                    )
-                    gr.Markdown("Note: The model shown here has a darker appearance. Download to get correct results.")
-
-            with gr.Row():
-                gr.Markdown('''Try a different <b>seed value</b> if the result is unsatisfying (Default: 42).''')
-
-    gr.Markdown(_CITE_)
-
+    
+    # Set up event handlers
     mv_images = gr.State()
-
-    submit.click(fn=check_input_image, inputs=[input_image]).success(
+    
+    # 3D Model generation flow
+    submit_3d.click(fn=check_input_image, inputs=[input_image]).success(
         fn=preprocess,
         inputs=[input_image, do_remove_background],
         outputs=[processed_image],
@@ -376,11 +376,21 @@ with gr.Blocks() as demo:
         fn=generate_mvs,
         inputs=[processed_image, sample_steps, sample_seed],
         outputs=[mv_images, mv_show_images]
-        
     ).success(
         fn=make3d,
         inputs=[mv_images],
         outputs=[output_model_obj, output_model_glb]
     )
+    
+    # Pixel Art generation flow
+    submit_pixel.click(
+        fn=generate_pixel_art,
+        inputs=[pixel_prompt, pixel_bg_remove, pixel_steps, pixel_seed],
+        outputs=[pixel_output]
+    )
+    
+    # Enable API access
+    demo.queue(concurrency_count=1)
 
-demo.launch()
+# Launch with API access enabled
+demo.launch(enable_api=True, share=False)