Update app.py

app.py

@@ -1,26 +1,35 @@
 #!/usr/bin/env python
 # -*- coding:utf-8 -*-
-# Power by Zongsheng Yue 2024-12-11 17:17:41
-# Modified to run on CPU
 
 import os
-# Force PyTorch to ignore GPUs and run on CPU
-os.environ["CUDA_VISIBLE_DEVICES"] = ""
+os.environ["CUDA_VISIBLE_DEVICES"] = ""  # 🔥 Force CPU
 
+import spaces
 import warnings
 warnings.filterwarnings("ignore")
+
 import argparse
 import numpy as np
 import gradio as gr
 from pathlib import Path
 from omegaconf import OmegaConf
 from sampler_invsr import InvSamplerSR
+
 from utils import util_common
 from utils import util_image
 from basicsr.utils.download_util import load_file_from_url
 
+# Optional: enforce CPU in torch if used internally
+try:
+    import torch
+    torch.set_default_device("cpu")
+except:
+    pass
+
+
 def get_configs(num_steps=1, chopping_size=128, seed=12345):
     configs = OmegaConf.load("./configs/sample-sd-turbo.yaml")
+
     if num_steps == 1:
         configs.timesteps = [200,]
     elif num_steps == 2:
@@ -36,12 +45,13 @@ def get_configs(num_steps=1, chopping_size=128, seed=12345):
         configs.timesteps = np.linspace(
             start=250, stop=0, num=num_steps, endpoint=False, dtype=np.int64()
         ).tolist()
+
     print(f'Setting timesteps for inference: {configs.timesteps}')
-    
-    # path to save noise predictor
+
     started_ckpt_name = "noise_predictor_sd_turbo_v5.pth"
     started_ckpt_dir = "./weights"
     util_common.mkdir(started_ckpt_dir, delete=False, parents=True)
+
     started_ckpt_path = Path(started_ckpt_dir) / started_ckpt_name
     if not started_ckpt_path.exists():
         load_file_from_url(
@@ -50,100 +60,63 @@ def get_configs(num_steps=1, chopping_size=128, seed=12345):
             progress=True,
             file_name=started_ckpt_name,
         )
+
     configs.model_start.ckpt_path = str(started_ckpt_path)
     configs.bs = 1
-    configs.seed = seed
+    configs.seed = seed 
     configs.basesr.chopping.pch_size = chopping_size
+
     if chopping_size == 128:
         configs.basesr.chopping.extra_bs = 8
     elif chopping_size == 256:
         configs.basesr.chopping.extra_bs = 4
     else:
         configs.basesr.chopping.extra_bs = 1
-        
-    # Explicitly set device to CPU if the config schema supports it
-    if OmegaConf.is_struct(configs):
-        OmegaConf.update(configs, 'device', 'cpu')
-    else:
-        configs.device = 'cpu'
-        
+
     return configs
 
-# Removed @spaces.GPU decorator for CPU execution
+
+# ❌ Removed @spaces.GPU
 def predict(in_path, num_steps=1, chopping_size=128, seed=12345):
     configs = get_configs(num_steps=num_steps, chopping_size=chopping_size, seed=seed)
+
     sampler = InvSamplerSR(configs)
+
     out_dir = Path('invsr_output')
-    if not out_dir.exists():
-        out_dir.mkdir()
+    out_dir.mkdir(exist_ok=True)
+
     sampler.inference(in_path, out_path=out_dir, bs=1)
+
     out_path = out_dir / f"{Path(in_path).stem}.png"
     assert out_path.exists(), 'Super-resolution failed!'
+
     im_sr = util_image.imread(out_path, chn="rgb", dtype="uint8")
+
     return im_sr, str(out_path)
 
+
 title = "Arbitrary-steps Image Super-resolution via Diffusion Inversion"
-description = r"""
-<b>Official Gradio demo</b> for <a href='https://github.com/zsyOAOA/InvSR' target='_blank'><b>Arbitrary-steps Image Super-resolution via Diffuion Inversion</b></a>.<br>
-🔥 InvSR is an image super-resolution method via Diffusion Inversion, supporting arbitrary sampling steps.<br>
-<i>⚠️ This instance is running on <b>CPU</b>. Inference will be significantly slower than on GPU. Consider using a lower chopping size and fewer steps for faster results.</i>
-"""
-article = r"""
-If you've found InvSR useful for your research or projects, please show your support by ⭐ the <a href='https://github.com/zsyOAOA/InvSR' target='_blank'>Github Repo</a>. Thanks!
-[![GitHub Stars](https://img.shields.io/github/stars/zsyOAOA/InvSR?affiliations=OWNER&color=green&style=social)](https://github.com/zsyOAOA/InvSR)
----
-If our work is useful for your research, please consider citing:
-```bibtex
-@article{yue2024InvSR,
-title={Arbitrary-steps Image Super-resolution via Diffusion Inversion},
-author={Yue, Zongsheng and Kang, Liao and Loy, Chen Change},
-journal = {arXiv preprint arXiv:2412.09013},
-year={2024},
-}
-```
-
-📋 **License**
-
-This project is licensed under <a rel="license" href="https://github.com/zsyOAOA/InvSR/blob/master/LICENSE">S-Lab License 1.0</a>.
-Redistribution and use for non-commercial purposes should follow this license.
-
-📧 **Contact**
-
-If you have any questions, please feel free to contact me via <b>zsyzam@gmail.com</b>.
-![visitors](https://visitor-badge.laobi.icu/badge?page_id=zsyOAOA/InvSR)
+
+description = """
+<b>CPU version</b> of InvSR demo.<br>
+⚠️ Note: Running on CPU will be significantly slower than GPU.
 """
+
 demo = gr.Interface(
     fn=predict,
     inputs=[
-        gr.Image(type="filepath", label="Input: Low Quality Image"),
-        gr.Dropdown(
-            choices=[1,2,3,4,5],
-            value=1,
-            label="Number of steps",
-            ),
-        gr.Dropdown(
-            choices=[128, 256, 512],
-            value=128,
-            label="Chopping size",
-            ),
-        gr.Number(value=12345, precision=0, label="Ranom seed")
+        gr.Image(type="filepath", label="Input Image"),
+        gr.Dropdown([1,2,3,4,5], value=1, label="Steps"),
+        gr.Dropdown([128, 256, 512], value=128, label="Chopping size"),
+        gr.Number(value=12345, precision=0, label="Seed")
     ],
     outputs=[
-        gr.Image(type="numpy", label="Output: High Quality Image"),
-        gr.File(label="Download the output")
+        gr.Image(type="numpy", label="Output Image"),
+        gr.File(label="Download")
     ],
     title=title,
-    description=description,
-    article=article,
-    examples=[
-        ['./testdata/RealSet80/29.jpg',  3, 128, 12345],
-        ['./testdata/RealSet80/32.jpg',  1, 128, 12345],
-        ['./testdata/RealSet80/0030.jpg',  1, 128, 12345],
-        ['./testdata/RealSet80/2684538-PH.jpg', 1, 128, 12345],
-        ['./testdata/RealSet80/oldphoto6.png', 1, 128, 12345],
-      ]
-    )
+    description=description
+)
 
 demo.queue(max_size=5)
-demo.launch(share=False)
-
+demo.launch()