Update app.py

app.py

@@ -17,7 +17,6 @@ subprocess.run(
     shell=True,
 )
 
-
 os.makedirs("/home/user/app/checkpoints", exist_ok=True)
 from huggingface_hub import snapshot_download
 snapshot_download(
@@ -51,7 +50,6 @@ import models
 from transport import Sampler, create_transport
 
 from multiprocessing import Process,Queue,set_start_method,get_context
-#set_start_method('fork')
 
 class ModelFailure:
     pass
@@ -110,11 +108,6 @@ def model_main(args, master_port, rank):
     # 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)
-
 
     train_args = torch.load(os.path.join(args.ckpt, "model_args.pth"))
     print("Loaded model arguments:", json.dumps(train_args.__dict__, indent=2))
@@ -307,16 +300,13 @@ def inference(args, infer_args, text_encoder, tokenizer, vae, model):
                 samples = (samples + 1.0) / 2.0
                 samples.clamp_(0.0, 1.0)
 
-                img = to_pil_image(samples[0, :].float())
+                img = to_pil_image(samples[0].float())
                 print("> generated image, done.")
 
-                # if response_queue is not None:
-                #     response_queue.put((img, metadata))
                 return img, metadata
             except Exception:
                 print(traceback.format_exc())
                 return ModelFailure()
-                # response_queue.put(ModelFailure())
 
 
 def none_or_str(value):
@@ -402,38 +392,11 @@ def main():
         raise NotImplementedError("Multi-GPU Inference is not yet supported")
 
     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)
 
     text_encoder, tokenizer, vae, model = model_main(args, master_port, 0)
 
-    description = args.ckpt.split('/')[-1]
-    #"""
-    # Lumina Next Text-to-Image
-
-    #Lumina-Next-T2I is a 2B Next-DiT model with 2B text encoder.
-
-    #Demo current model: `Lumina-Next-T2I`
+    description = "Lumina-Image 2.0 ([Github](https://github.com/Alpha-VLLM/Lumina-Image-2.0/tree/main))"
 
-    #"""
     with gr.Blocks() as demo:
         with gr.Row():
             gr.Markdown(description)
@@ -568,13 +531,7 @@ def main():
             )  # noqa
 
         @spaces.GPU(duration=200)
-        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
+        def on_submit(*infer_args, progress=gr.Progress(track_tqdm=True),):
             result = inference(args, infer_args, text_encoder, tokenizer, vae, model)
             if isinstance(result, ModelFailure):
                 raise RuntimeError("Model failed to generate the image.")