Update app.py

app.py

@@ -1,6 +1,6 @@
-
 import gradio as gr
 from PIL import Image
+import requests
 
 from tld.denoiser import Denoiser
 from tld.diffusion import DiffusionGenerator
@@ -19,10 +19,13 @@ device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 to_pil = transforms.ToPILImage()
 
 
-###config:
-vae_scale_factor = 8
-img_size = 32
-model_dtype = torch.float32
+def download_file(url, filename):
+
+    with requests.get(url, stream=True) as r:
+        r.raise_for_status() 
+        with open(filename, 'wb') as f:
+            for chunk in r.iter_content(chunk_size=8192): 
+                f.write(chunk)
 
 @torch.no_grad()
 def encode_text(label, model):
@@ -55,6 +58,14 @@ def generate_image_from_text(prompt, class_guidance=6, seed=11, num_imgs=1, img_
     print("Images Generated and Saved. They will shortly output below.")
     return out
 
+###config:
+vae_scale_factor = 8
+img_size = 32
+model_dtype = torch.float32
+
+file_url = "https://huggingface.co/apapiu/small_ldt/resolve/main/state_dict_378000.pth"
+local_filename = "state_dict_378000.pth"
+download_file(file_url, local_filename)
 
 
 denoiser = Denoiser(image_size=32, noise_embed_dims=256, patch_size=2,
@@ -85,4 +96,4 @@ iface = gr.Interface(
 )
 
 # Launch the app
-iface.launch()
+iface.launch()