Update app.py

app.py

@@ -70,8 +70,6 @@ def encode_prompt(prompt_batch, text_encoder, tokenizer, proportion_empty_prompt
             truncation=True,
             return_tensors="pt",
         )
-        for name, param in text_encoder.named_parameters():
-            print(name, param.device)
 
         print(f"Text Encoder Device: {text_encoder.device}")
         text_input_ids = text_inputs.input_ids.cuda()
@@ -90,7 +88,7 @@ def encode_prompt(prompt_batch, text_encoder, tokenizer, proportion_empty_prompt
 
 
 @torch.no_grad()
-def model_main(args, master_port, rank, request_queue, response_queue, mp_barrier):
+def model_main(args, master_port, rank):
     # import here to avoid huggingface Tokenizer parallelism warnings
     from diffusers.models import AutoencoderKL
     from transformers import AutoModel, AutoTokenizer
@@ -106,10 +104,10 @@ def model_main(args, master_port, rank, request_queue, response_queue, mp_barrie
     # Override the built-in print with the new version
     builtins.print = print
 
-    os.environ["MASTER_PORT"] = str(master_port)
-    os.environ["MASTER_ADDR"] = "127.0.0.1"
-    os.environ["RANK"] = str(rank)
-    os.environ["WORLD_SIZE"] = str(args.num_gpus)
+    # os.environ["MASTER_PORT"] = str(master_port)
+    # os.environ["MASTER_ADDR"] = "127.0.0.1"
+    # os.environ["RANK"] = str(rank)
+    # os.environ["WORLD_SIZE"] = str(args.num_gpus)
 
 
     train_args = torch.load(os.path.join(args.ckpt, "model_args.pth"))
@@ -159,8 +157,12 @@ def model_main(args, master_port, rank, request_queue, response_queue, mp_barrie
         ckpt = torch.load(ckpt_path, map_location="cuda")
     model.load_state_dict(ckpt, strict=True)
     print('load model finish')
-    mp_barrier.wait()
+    
+    return text_encoder, tokenizer, vae, model
+
 
+@torch.no_grad()
+def inference(args, infer_args, text_encoder, tokenizer, vae, model)
     with torch.autocast("cuda", dtype):
         while True:
             (
@@ -178,7 +180,7 @@ def model_main(args, master_port, rank, request_queue, response_queue, mp_barrie
                 scaling_method,
                 scaling_watershed,
                 proportional_attn,
-            ) = request_queue.get()
+            ) = infer_args
 
            
             system_prompt = system_type
@@ -243,13 +245,13 @@ def model_main(args, master_port, rank, request_queue, response_queue, mp_barrie
                     torch.random.manual_seed(int(seed))
                 z = torch.randn([1, 16, latent_h, latent_w], device="cuda").to(dtype)
                 z = z.repeat(2, 1, 1, 1)
-                model.cpu()
+
                 with torch.no_grad():
                     if neg_cap != "":
                         cap_feats, cap_mask = encode_prompt([cap] + [neg_cap], text_encoder, tokenizer, 0.0)
                     else:
                         cap_feats, cap_mask = encode_prompt([cap] + [""], text_encoder, tokenizer, 0.0)
-                model.cuda()
+               
                 cap_mask = cap_mask.to(cap_feats.device)
 
                 model_kwargs = dict(
@@ -297,12 +299,13 @@ def model_main(args, master_port, rank, request_queue, response_queue, mp_barrie
                 img = to_pil_image(samples[0, :].float())
                 print("> generated image, done.")
 
-                if response_queue is not None:
-                    response_queue.put((img, metadata))
-
+                # if response_queue is not None:
+                #     response_queue.put((img, metadata))
+                return img, metadata
             except Exception:
                 print(traceback.format_exc())
-                response_queue.put(ModelFailure())
+                return ModelFailure()
+                # response_queue.put(ModelFailure())
 
 
 def none_or_str(value):
@@ -389,25 +392,27 @@ def main():
 
     master_port = find_free_port()
     #mp.set_start_method("fork")
-    processes = []
-    request_queues = []
-    response_queue = mp.Queue()
-    mp_barrier = mp.Barrier(args.num_gpus + 1)
-    for i in range(args.num_gpus):
-        request_queues.append(mp.Queue())
-        p = mp.Process(
-            target=model_main,
-            args=(
-                args,
-                master_port,
-                i,
-                request_queues[i],
-                response_queue if i == 0 else None,
-                mp_barrier,
-            ),
-        )
-        p.start()
-        processes.append(p)
+    # processes = []
+    # request_queues = []
+    # response_queue = mp.Queue()
+    # mp_barrier = mp.Barrier(args.num_gpus + 1)
+    # for i in range(args.num_gpus):
+    #     request_queues.append(mp.Queue())
+    #     p = mp.Process(
+    #         target=model_main,
+    #         args=(
+    #             args,
+    #             master_port,
+    #             i,
+    #             request_queues[i],
+    #             response_queue if i == 0 else None,
+    #             mp_barrier,
+    #         ),
+    #     )
+    #     p.start()
+    #     processes.append(p)
+
+    model_main(args, master_port, 0)
 
     description = args.ckpt.split('/')[-1]
     #"""
@@ -552,15 +557,18 @@ def main():
             )  # noqa
 
         @spaces.GPU(duration=200)
-        def on_submit(*args):
-            for q in request_queues:
-                q.put(args)
-            result = response_queue.get()
+        def on_submit(*infer_args):
+            # for q in request_queues:
+            #     q.put(args)
+            # result = response_queue.get()
+            # if isinstance(result, ModelFailure):
+            #     raise RuntimeError
+            # img, metadata = result
+            result = inference(args, infer_args, text_encoder, tokenizer, vae, model)
             if isinstance(result, ModelFailure):
-                raise RuntimeError
-            img, metadata = result
+                raise RuntimeError("Model failed to generate the image.")
 
-            return img, metadata
+            return result
 
         submit_btn.click(
             on_submit,