back

app.py

@@ -26,54 +26,53 @@ def setup():
     os.system("pip install tokenizers x-transformers==0.22.0 axial-positional-embedding")
     os.system("pip install einops==0.3.2")
     sys.path.append("improved-diffusion")
-    from improved_diffusion import pipeline
-    # import improved_diffusion.pipeline
-    # from transformer_utils.util.tfm_utils import get_local_path_from_huggingface_cdn
-    #
-    # if not os.path.exists(model_path_diffusion):
-    #     model_tar_name = 'model.tar'
-    #     model_tar_path = get_local_path_from_huggingface_cdn(
-    #         HF_REPO_NAME_DIFFUSION, model_tar_name
-    #     )
-    #     subprocess.run(f"tar -xf {model_tar_path} && rm {model_tar_path}", shell=True)
-    #
-    # checkpoint_path_sres1 = os.path.join(model_path_diffusion, "sres1.pt")
-    # config_path_sres1 = os.path.join(model_path_diffusion, "config_sres1.json")
-    #
-    # checkpoint_path_sres2 = os.path.join(model_path_diffusion, "sres2.pt")
-    # config_path_sres2 = os.path.join(model_path_diffusion, "config_sres2.json")
-    #
-    # # load
-    # sampling_model_sres1 = improved_diffusion.pipeline.SamplingModel.from_config(
-    #     checkpoint_path=checkpoint_path_sres1,
-    #     config_path=config_path_sres1,
-    #     timestep_respacing=timestep_respacing_sres1
-    # )
-    #
-    # sampling_model_sres2 = improved_diffusion.pipeline.SamplingModel.from_config(
-    #     checkpoint_path=checkpoint_path_sres2,
-    #     config_path=config_path_sres2,
-    #     timestep_respacing=timestep_respacing_sres2
-    # )
-    #
-    # pipeline = improved_diffusion.pipeline.SamplingPipeline(sampling_model_sres1, sampling_model_sres2)
+
+    import improved_diffusion.pipeline
+    from transformer_utils.util.tfm_utils import get_local_path_from_huggingface_cdn
+
+    if not os.path.exists(model_path_diffusion):
+        model_tar_name = 'model.tar'
+        model_tar_path = get_local_path_from_huggingface_cdn(
+            HF_REPO_NAME_DIFFUSION, model_tar_name
+        )
+        subprocess.run(f"tar -xf {model_tar_path} && rm {model_tar_path}", shell=True)
+
+    checkpoint_path_sres1 = os.path.join(model_path_diffusion, "sres1.pt")
+    config_path_sres1 = os.path.join(model_path_diffusion, "config_sres1.json")
+
+    checkpoint_path_sres2 = os.path.join(model_path_diffusion, "sres2.pt")
+    config_path_sres2 = os.path.join(model_path_diffusion, "config_sres2.json")
+
+    # load
+    sampling_model_sres1 = improved_diffusion.pipeline.SamplingModel.from_config(
+        checkpoint_path=checkpoint_path_sres1,
+        config_path=config_path_sres1,
+        timestep_respacing=timestep_respacing_sres1
+    )
+
+    sampling_model_sres2 = improved_diffusion.pipeline.SamplingModel.from_config(
+        checkpoint_path=checkpoint_path_sres2,
+        config_path=config_path_sres2,
+        timestep_respacing=timestep_respacing_sres2
+    )
+
+    pipeline = improved_diffusion.pipeline.SamplingPipeline(sampling_model_sres1, sampling_model_sres2)
     return pipeline
 
 pipeline = setup()
 
 def handler(text, ts1, ts2, gs1):
-    pass
-    # # a = np.random.randint(0, 255, (128, 128, 3)).astype(np.uint8)
-    # data = {'text': text[:380], 'guidance_scale': gs1}
-    # args = {k: v for k, v in DIFFUSION_DEFAULTS.items()}
-    # args.update(data)
-    #
-    # print(f"running: {args}")
-    #
-    # pipeline.base_model.set_timestep_respacing(str(ts1))
-    # pipeline.super_res_model.set_timestep_respacing(str(ts2))
-    #
-    # return pipeline.sample(**args)
+    # a = np.random.randint(0, 255, (128, 128, 3)).astype(np.uint8)
+    data = {'text': text[:380], 'guidance_scale': gs1}
+    args = {k: v for k, v in DIFFUSION_DEFAULTS.items()}
+    args.update(data)
+
+    print(f"running: {args}")
+
+    pipeline.base_model.set_timestep_respacing(str(ts1))
+    pipeline.super_res_model.set_timestep_respacing(str(ts2))
+
+    return pipeline.sample(**args)
 
 
 text = st.text_area('asdf')