Update app.py

app.py

@@ -1,8 +1,3 @@
-"""
-This file is used for deploying hugging face demo:
-https://huggingface.co/spaces/
-"""
-
 import sys
 sys.path.append('StableSR')
 import os
@@ -25,48 +20,41 @@ from scripts.wavelet_color_fix import wavelet_reconstruction, adaptive_instance_
 from scripts.util_image import ImageSpliterTh
 from basicsr.utils.download_util import load_file_from_url
 from einops import rearrange, repeat
+from itertools import islice
 
-# os.system("pip freeze")
-
+# Download weights
 pretrain_model_url = {
-	'stablesr_512': 'https://huggingface.co/Iceclear/StableSR/resolve/main/stablesr_000117.ckpt',
-	'stablesr_768': 'https://huggingface.co/Iceclear/StableSR/resolve/main/stablesr_768v_000139.ckpt',
-	'CFW': 'https://huggingface.co/Iceclear/StableSR/resolve/main/vqgan_cfw_00011.ckpt',
+    'stablesr_512': 'https://huggingface.co/Iceclear/StableSR/resolve/main/stablesr_000117.ckpt',
+    'stablesr_768': 'https://huggingface.co/Iceclear/StableSR/resolve/main/stablesr_768v_000139.ckpt',
+    'CFW': 'https://huggingface.co/Iceclear/StableSR/resolve/main/vqgan_cfw_00011.ckpt',
 }
-# download weights
-if not os.path.exists('./stablesr_000117.ckpt'):
-	load_file_from_url(url=pretrain_model_url['stablesr_512'], model_dir='./', progress=True, file_name=None)
-if not os.path.exists('./stablesr_768v_000139.ckpt'):
-	load_file_from_url(url=pretrain_model_url['stablesr_768'], model_dir='./', progress=True, file_name=None)
-if not os.path.exists('./vqgan_cfw_00011.ckpt'):
-	load_file_from_url(url=pretrain_model_url['CFW'], model_dir='./', progress=True, file_name=None)
-
-# download images
-torch.hub.download_url_to_file(
-	'https://raw.githubusercontent.com/zsyOAOA/ResShift/master/testdata/RealSet128/Lincoln.png',
-	'01.png')
-torch.hub.download_url_to_file(
-	'https://raw.githubusercontent.com/zsyOAOA/ResShift/master/testdata/RealSet128/oldphoto6.png',
-	'02.png')
-torch.hub.download_url_to_file(
-	'https://raw.githubusercontent.com/zsyOAOA/ResShift/master/testdata/RealSet128/comic2.png',
-	'03.png')
-torch.hub.download_url_to_file(
-	'https://raw.githubusercontent.com/zsyOAOA/ResShift/master/testdata/RealSet128/OST_120.png',
-	'04.png')
-torch.hub.download_url_to_file(
-	'https://raw.githubusercontent.com/zsyOAOA/ResShift/master/testdata/RealSet65/comic3.png',
-	'05.png')
+
+for k, url in pretrain_model_url.items():
+    filename = url.split("/")[-1]
+    if not os.path.exists(f'./{filename}'):
+        load_file_from_url(url=url, model_dir='./', progress=True, file_name=None)
+
+# Download sample images
+image_urls = [
+    ('01.png', 'https://raw.githubusercontent.com/zsyOAOA/ResShift/master/testdata/RealSet128/Lincoln.png'),
+    ('02.png', 'https://raw.githubusercontent.com/zsyOAOA/ResShift/master/testdata/RealSet128/oldphoto6.png'),
+    ('03.png', 'https://raw.githubusercontent.com/zsyOAOA/ResShift/master/testdata/RealSet128/comic2.png'),
+    ('04.png', 'https://raw.githubusercontent.com/zsyOAOA/ResShift/master/testdata/RealSet128/OST_120.png'),
+    ('05.png', 'https://raw.githubusercontent.com/zsyOAOA/ResShift/master/testdata/RealSet65/comic3.png'),
+]
+
+for fname, url in image_urls:
+    torch.hub.download_url_to_file(url, fname)
 
 def load_img(path):
-	image = Image.open(path).convert("RGB")
-	w, h = image.size
-	w, h = map(lambda x: x - x % 32, (w, h))  # resize to integer multiple of 32
-	image = image.resize((w, h), resample=PIL.Image.LANCZOS)
-	image = np.array(image).astype(np.float32) / 255.0
-	image = image[None].transpose(0, 3, 1, 2)
-	image = torch.from_numpy(image)
-	return 2.*image - 1.
+    image = Image.open(path).convert("RGB")
+    w, h = image.size
+    w, h = map(lambda x: x - x % 32, (w, h))
+    image = image.resize((w, h), resample=PIL.Image.LANCZOS)
+    image = np.array(image).astype(np.float32) / 255.0
+    image = image[None].transpose(0, 3, 1, 2)
+    image = torch.from_numpy(image)
+    return 2.*image - 1.
 
 def space_timesteps(num_timesteps, section_counts):
 	"""
@@ -143,11 +131,13 @@ def load_model_from_config(config, ckpt, verbose=False):
 	model.eval()
 	return model
 
-# device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+# Load VQGAN model
 device = torch.device("cuda")
 vqgan_config = OmegaConf.load("StableSR/configs/autoencoder/autoencoder_kl_64x64x4_resi.yaml")
-vq_model = load_model_from_config(vqgan_config, './vqgan_cfw_00011.ckpt')
-vq_model = vq_model.to(device)
+vq_model = instantiate_from_config(vqgan_config.model)
+vq_sd = torch.load('./vqgan_cfw_00011.ckpt', map_location='cpu')['state_dict']
+vq_model.load_state_dict(vq_sd, strict=False)
+vq_model.cuda().eval()
 
 os.makedirs('output', exist_ok=True)
 
@@ -284,6 +274,7 @@ def inference(image, upscale, dec_w, seed, model_type, ddpm_steps, colorfix_type
 		print('Global exception', error)
 		return None, None
 
+# Gradio UI
 with gr.Blocks(title="Exploiting Diffusion Prior for Real-World Image Super-Resolution") as demo:
     gr.Markdown(
         """
@@ -298,7 +289,7 @@ with gr.Blocks(title="Exploiting Diffusion Prior for Real-World Image Super-Reso
         If StableSR is helpful, please help to ⭐ the <a href='https://github.com/IceClear/StableSR' target='_blank'>Github Repo</a>. Thanks!  
         [![GitHub Stars](https://img.shields.io/github/stars/IceClear/StableSR?style=social)](https://github.com/IceClear/StableSR)  
         ---
-        
+
         📝 **Citation**  
         If our work is useful for your research, please consider citing:
 
@@ -319,59 +310,49 @@ with gr.Blocks(title="Exploiting Diffusion Prior for Real-World Image Super-Reso
         If you have any questions, please feel free to reach me out at <b>iceclearwjy@gmail.com</b>.
 
         <div>
-        	🤗 Find Me:  
-        	<a href="https://twitter.com/Iceclearwjy"><img style="margin-top:0.5em; margin-bottom:0.5em" src="https://img.shields.io/twitter/follow/Iceclearwjy?label=%40Iceclearwjy&style=social" alt="Twitter Follow"></a>  
-        	<a href="https://github.com/IceClear"><img style="margin-top:0.5em; margin-bottom:2em" src="https://img.shields.io/github/followers/IceClear?style=social" alt="Github Follow"></a>  
+        🤗 Find Me:  
+        <a href="https://twitter.com/Iceclearwjy"><img style="margin-top:0.5em; margin-bottom:0.5em" src="https://img.shields.io/twitter/follow/Iceclearwjy?label=%40Iceclearwjy&style=social" alt="Twitter Follow"></a>  
+        <a href="https://github.com/IceClear"><img style="margin-top:0.5em; margin-bottom:2em" src="https://img.shields.io/github/followers/IceClear?style=social" alt="Github Follow"></a>  
         </div>
 
         <center><img src='https://visitor-badge.laobi.icu/badge?page_id=IceClear/StableSR' alt='visitors'></center>
         """
     )
 
-	with gr.Row():
-		with gr.Column():
-			image = gr.Image(type="filepath", label="Input")
-			upscale = gr.Number(value=1, label="Rescaling_Factor (Large images require huge time)")
-			dec_w = gr.Slider(0, 1, value=0.5, step=0.01, label='CFW_Fidelity (0 for better quality, 1 for better identity)')
-			seed = gr.Number(value=42, label="Seeds")
-			model_type = gr.Dropdown(
-				choices=["512", "768v"],
-				value="512",
-				label="Model"
-			)
-			ddpm_steps = gr.Slider(10, 1000, value=200, step=1, label='Sampling timesteps for DDPM')
-			colorfix_type = gr.Dropdown(
-				choices=["none", "adain", "wavelet"],
-				value="adain",
-				label="Color_Correction"
-			)
-			run_btn = gr.Button("Run Inference")
-
-		with gr.Column():
-			output_image = gr.Image(type="numpy", label="Output")
-			output_file = gr.File(label="Download the output")
-
-	# Inference trigger
-	run_btn.click(
-		fn=inference,
-		inputs=[image, upscale, dec_w, seed, model_type, ddpm_steps, colorfix_type],
-		outputs=[output_image, output_file]
-	)
-
-	# Example section
-	gr.Examples(
-		examples=[
-			['./01.png', 4, 0.5, 42, "512", 200, "adain"],
-			['./02.png', 4, 0.5, 42, "512", 200, "adain"],
-			['./03.png', 4, 0.5, 42, "512", 200, "adain"],
-			['./04.png', 4, 0.5, 42, "512", 200, "adain"],
-			['./05.png', 4, 0.5, 42, "512", 200, "adain"]
-		],
-		fn=inference,
-		inputs=[image, upscale, dec_w, seed, model_type, ddpm_steps, colorfix_type],
-		outputs=[output_image, output_file],
-		cache_examples=True
-	)
+    with gr.Row():
+        with gr.Column():
+            image = gr.Image(type="filepath", label="Input")
+            upscale = gr.Number(value=1, label="Rescaling_Factor")
+            dec_w = gr.Slider(0, 1, value=0.5, step=0.01, label='CFW_Fidelity')
+            seed = gr.Number(value=42, label="Seeds")
+            model_type = gr.Dropdown(choices=["512", "768v"], value="512", label="Model")
+            ddpm_steps = gr.Slider(10, 1000, value=200, step=1, label='DDPM Steps')
+            colorfix_type = gr.Dropdown(choices=["none", "adain", "wavelet"], value="adain", label="Color Correction")
+            run_btn = gr.Button("Run Inference")
+
+        with gr.Column():
+            output_image = gr.Image(type="numpy", label="Output")
+            output_file = gr.File(label="Download the output")
+
+    run_btn.click(
+        fn=inference,
+        inputs=[image, upscale, dec_w, seed, model_type, ddpm_steps, colorfix_type],
+        outputs=[output_image, output_file]
+    )
+
+    gr.Examples(
+        examples=[
+            ['01.png', 4, 0.5, 42, "512", 200, "adain"],
+            ['02.png', 4, 0.5, 42, "512", 200, "adain"],
+            ['03.png', 4, 0.5, 42, "512", 200, "adain"],
+            ['04.png', 4, 0.5, 42, "512", 200, "adain"],
+            ['05.png', 4, 0.5, 42, "512", 200, "adain"]
+        ],
+        fn=inference,
+        inputs=[image, upscale, dec_w, seed, model_type, ddpm_steps, colorfix_type],
+        outputs=[output_image, output_file],
+        cache_examples=True
+    )
 
 demo.queue()
-demo.launch()
+demo.launch()