Update app.py

app.py

@@ -12,12 +12,14 @@ from diffusion_model.sampler.ddim import DDIM
 from diffusion_model.models.latent_diffusion_model import LatentDiffusionModel
 from diffusion_model.network.unet import Unet
 from diffusion_model.network.unet_wrapper import UnetWrapper
+from huggingface_hub import hf_hub_download
 
 # import spaces #[uncomment to use ZeroGPU]
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 loader = Loader(device)
 repo_id = "JuyeopDang/KoFace-Diffusion"
+CONFIG_PATH = 'configs/composite_config.yaml'
 
 if torch.cuda.is_available():
     torch_dtype = torch.float16
@@ -34,9 +36,6 @@ def load_model_from_HF(model, repo_id, filename, is_ema=False):
     return model
 
 if __name__ == "__main__":
-    from huggingface_hub import hf_hub_download
-    CONFIG_PATH = 'configs/composite_config.yaml'
-
     vae = VariationalAutoEncoder(CONFIG_PATH)
     sampler = DDIM(CONFIG_PATH)
     clip = KoCLIPWrapper()
@@ -47,5 +46,21 @@ if __name__ == "__main__":
     vae = load_model_from_HF(vae, repo_id, "composite_epoch2472.pth", False)
     clip = load_model_from_HF(clip, repo_id, "asian-composite-fine-tuned-koclip.pth", True)
     dm = load_model_from_HF(dm, repo_id, "asian-composite-clip-ldm.pth", True)
+
+    def generate_image(y, gamma, dm):
+        images = dm.sample(2, y = y, gamma = gamma)
+        images = images.permute(0, 2, 3, 1) 
+        if type(images) is torch.Tensor:
+          images = images.detach().cpu().numpy()
+        images = np.clip(images / 2 + 0.5, 0, 1)
+        return im.fromarray((images[0] * 255).astype(np.uint8))
+
+    demo = gr.Interface(
+        generate_image,
+        inputs=["textbox", gr.Slider(0, 10)],
+        outputs=["image"],
+    )
     
-    print(dm)
+    demo.launch()
+        
+