Enhance gen_image function in app.py to return base64 encoded GLB data alongside generated images. Clean up argument parsing by removing unused stage1 and stage2 config options in parser. Update requirements.txt to specify gradio version for compatibility.

app.py

@@ -93,7 +93,9 @@ def preprocess_image(image, background_choice, foreground_ratio, backgroud_color
     image = add_background(image, backgroud_color)
     return image.convert("RGB")
 
+
 @spaces.GPU
+
 def gen_image(input_image, seed, scale, step):
     global pipeline, model, args
     pipeline.set_seed(seed)
@@ -102,24 +104,19 @@ def gen_image(input_image, seed, scale, step):
     stage2_images = rt_dict["stage2_images"]
     np_imgs = np.concatenate(stage1_images, 1)
     np_xyzs = np.concatenate(stage2_images, 1)
-
+    
     glb_path = generate3d(model, np_imgs, np_xyzs, args.device)
-    return Image.fromarray(np_imgs), Image.fromarray(np_xyzs), glb_path#, obj_path
+    
+    # Read the GLB file and encode it in base64
+    with open(glb_path, 'rb') as f:
+        glb_bytes = f.read()
+    encoded_glb = 'data:model/gltf-binary;base64,' + base64.b64encode(glb_bytes).decode('utf-8')
+    
+    # Return images and the encoded GLB data
+    return Image.fromarray(np_imgs), Image.fromarray(np_xyzs), encoded_glb
 
 
 parser = argparse.ArgumentParser()
-parser.add_argument(
-    "--stage1_config",
-    type=str,
-    default="configs/nf7_v3_SNR_rd_size_stroke.yaml",
-    help="config for stage1",
-)
-parser.add_argument(
-    "--stage2_config",
-    type=str,
-    default="configs/stage2-v2-snr.yaml",
-    help="config for stage2",
-)
 
 parser.add_argument("--device", type=str, default="cuda")
 args = parser.parse_args()

requirements.txt

@@ -1,5 +1,4 @@
-gradio
-huggingface-hub
+huggingface-hub == 0.19.4
 diffusers==0.24.0
 einops==0.7.0
 Pillow==10.1.0
@@ -8,6 +7,7 @@ open-clip-torch==2.7.0
 opencv-contrib-python-headless==4.9.0.80
 opencv-python-headless==4.9.0.80
 xformers
+gradio==0.16.0
 omegaconf
 rembg
 git+https://github.com/NVlabs/nvdiffrast