queue 1

6DoF/dataset.py

@@ -2,40 +2,12 @@ import os
 import math
 from pathlib import Path
 import torch
-import torchvision
-from torch.utils.data import Dataset, DataLoader
-from torchvision import transforms
+from torch.utils.data import Dataset
 from PIL import Image
 import numpy as np
-import webdataset as wds
-from torch.utils.data.distributed import DistributedSampler
 import matplotlib.pyplot as plt
 import sys
 
-class ObjaverseDataLoader():
-    def __init__(self, root_dir, batch_size, total_view=12, num_workers=4):
-        self.root_dir = root_dir
-        self.batch_size = batch_size
-        self.num_workers = num_workers
-        self.total_view = total_view
-
-        image_transforms = [torchvision.transforms.Resize((256, 256)),
-                            transforms.ToTensor(),
-                            transforms.Normalize([0.5], [0.5])]
-        self.image_transforms = torchvision.transforms.Compose(image_transforms)
-
-    def train_dataloader(self):
-        dataset = ObjaverseData(root_dir=self.root_dir, total_view=self.total_view, validation=False,
-                                image_transforms=self.image_transforms)
-        # sampler = DistributedSampler(dataset)
-        return wds.WebLoader(dataset, batch_size=self.batch_size, num_workers=self.num_workers, shuffle=False)
-                             # sampler=sampler)
-
-    def val_dataloader(self):
-        dataset = ObjaverseData(root_dir=self.root_dir, total_view=self.total_view, validation=True,
-                                image_transforms=self.image_transforms)
-        sampler = DistributedSampler(dataset)
-        return wds.WebLoader(dataset, batch_size=self.batch_size, num_workers=self.num_workers, shuffle=False)
 
 def get_pose(transformation):
     # transformation: 4x4

app.py

@@ -183,19 +183,11 @@ def run_eschernet(eschernet_input_dict, sample_steps, sample_seed, nvs_num, nvs_
     # run inference
     # pipeline.to(device)
     pipeline.enable_xformers_memory_efficient_attention()
-    if CaPE_TYPE == "6DoF":
-        with torch.autocast("cuda"):
-            image = pipeline(input_imgs=input_image, prompt_imgs=input_image,
-                             poses=[[pose_out, pose_out_inv], [pose_in, pose_in_inv]],
-                             height=h, width=w, T_in=T_in, T_out=T_out,
-                             guidance_scale=guidance_scale, num_inference_steps=50, generator=generator,
-                             output_type="numpy").images
-    elif CaPE_TYPE == "4DoF":
-        with torch.autocast("cuda"):
-            image = pipeline(input_imgs=input_image, prompt_imgs=input_image, poses=[pose_out, pose_in],
-                             height=h, width=w, T_in=T_in, T_out=T_out,
-                             guidance_scale=guidance_scale, num_inference_steps=50, generator=generator,
-                             output_type="numpy").images
+    image = pipeline(input_imgs=input_image, prompt_imgs=input_image,
+                         poses=[[pose_out, pose_out_inv], [pose_in, pose_in_inv]],
+                         height=h, width=w, T_in=T_in, T_out=T_out,
+                         guidance_scale=guidance_scale, num_inference_steps=50, generator=generator,
+                         output_type="numpy").images
 
     # save output image
     output_dir = os.path.join(tmpdirname, "eschernet")
@@ -748,7 +740,7 @@ with gr.Blocks() as demo:
 
 # demo.queue(max_size=10)
 # demo.launch(share=True, server_name="0.0.0.0", server_port=None)
-demo.queue(max_size=1).launch()
+demo.queue(max_size=10).launch()
 
 # if __name__ == '__main__':
 #     main()