Update app.py

app.py

@@ -12,14 +12,17 @@ from huggingface_hub import hf_hub_download
 import subprocess
 
 # Install Real-ESRGAN with dependencies
-subprocess.run("pip install git+https://github.com/inference-sh/Real-ESRGAN.git basicsr opencv-python-headless", shell=True)
+try:
+    subprocess.run("pip install git+https://github.com/inference-sh/Real-ESRGAN.git basicsr opencv-python-headless --no-cache-dir", shell=True, check=True)
+except subprocess.CalledProcessError as e:
+    print(f"Failed to install Real-ESRGAN: {e}")
 
 from RealESRGAN import RealESRGAN
 
 # Force CPU usage
 device = torch.device("cpu")
-ENABLE_CPU_OFFLOAD = True  # Enable CPU offloading to manage memory
-USE_TORCH_COMPILE = False  # Disable torch.compile for CPU compatibility
+ENABLE_CPU_OFFLOAD = True
+USE_TORCH_COMPILE = False
 
 # Create model directories
 os.makedirs("models/Stable-diffusion", exist_ok=True)
@@ -27,21 +30,17 @@ os.makedirs("models/ControlNet", exist_ok=True)
 os.makedirs("models/VAE", exist_ok=True)
 os.makedirs("models/upscalers", exist_ok=True)
 
-# Download essential models (reduced set to save storage)
-def download_models():
-    models = {
-        "MODEL": ("dantea1118/juggernaut_reborn", "juggernaut_reborn.safetensors", "models/Stable-diffusion"),
-        "CONTROLNET": ("lllyasviel/ControlNet-v1-1", "control_v11f1e_sd15_tile.pth", "models/ControlNet"),
-        "VAE": ("stabilityai/sd-vae-ft-mse-original", "vae-ft-mse-840000-ema-pruned.safetensors", "models/VAE"),
-        "UPSCALER_X2": ("ai-forever/Real-ESRGAN", "RealESRGAN_x2.pth", "models/upscalers"),
-    }
-    for model, (repo_id, filename, local_dir) in models.items():
-        print(f"Downloading {model}...")
+# Download models on-demand
+def download_model(repo_id, filename, local_dir):
+    try:
+        print(f"Downloading {filename} from {repo_id}...")
         hf_hub_download(repo_id=repo_id, filename=filename, local_dir=local_dir)
+        print(f"Successfully downloaded {filename}")
+    except Exception as e:
+        print(f"Failed to download {filename}: {e}")
+        raise
 
-download_models()
-
-# Timer decorator for performance tracking
+# Timer decorator
 def timer_func(func):
     def wrapper(*args, **kwargs):
         start_time = time.time()
@@ -50,7 +49,7 @@ def timer_func(func):
         return result
     return wrapper
 
-# Lazy pipeline for memory efficiency
+# Lazy pipeline
 class LazyLoadPipeline:
     def __init__(self):
         self.pipe = None
@@ -59,6 +58,15 @@ class LazyLoadPipeline:
     def load(self):
         if self.pipe is None:
             print("Setting up pipeline...")
+            # Download models if not present
+            for model, (repo_id, filename, local_dir) in [
+                ("MODEL", ("dantea1118/juggernaut_reborn", "juggernaut_reborn.safetensors", "models/Stable-diffusion")),
+                ("CONTROLNET", ("lllyasviel/ControlNet-v1-1", "control_v11f1e_sd15_tile.pth", "models/ControlNet")),
+                ("VAE", ("stabilityai/sd-vae-ft-mse-original", "vae-ft-mse-840000-ema-pruned.safetensors", "models/VAE")),
+            ]:
+                if not os.path.exists(os.path.join(local_dir, filename)):
+                    download_model(repo_id, filename, local_dir)
+
             controlnet = ControlNetModel.from_single_file(
                 "models/ControlNet/control_v11f1e_sd15_tile.pth", torch_dtype=torch.float16
             )
@@ -79,15 +87,15 @@ class LazyLoadPipeline:
             if ENABLE_CPU_OFFLOAD:
                 print("Enabling CPU offloading...")
                 pipe.enable_model_cpu_offload()
-            return pipe
+            self.pipe = pipe
         return self.pipe
 
     def __call__(self, *args, **kwargs):
         if self.pipe is None:
-            self.pipe = self.load()
+            self.load()
         return self.pipe(*args, **kwargs)
 
-# Lazy Real-ESRGAN upscaler
+# Lazy Real-ESRGAN
 class LazyRealESRGAN:
     def __init__(self, device, scale):
         self.device = device
@@ -96,6 +104,8 @@ class LazyRealESRGAN:
 
     def load_model(self):
         if self.model is None:
+            if not os.path.exists(f"models/upscalers/RealESRGAN_x{self.scale}.pth"):
+                download_model("ai-forever/Real-ESRGAN", f"RealESRGAN_x{self.scale}.pth", "models/upscalers")
             self.model = RealESRGAN(self.device, scale=self.scale)
             self.model.load_weights(f'models/upscalers/RealESRGAN_x{self.scale}.pth', download=False)
 
@@ -121,7 +131,7 @@ def create_hdr_effect(original_image, hdr):
     if hdr == 0:
         return original_image
     cv_original = cv2.cvtColor(np.array(original_image), cv2.COLOR_RGB2BGR)
-    factors = [1.0 - 0.9 * hdr, 1.0 - 0.7 * hdr, 1.0, 1.0 + 0.2 * hdr]
+    factors = [1.0 - 0.7 * hdr, 1.0, 1.0 + 0.2 * hdr]
     images = [cv2.convertScaleAbs(cv_original, alpha=factor) for factor in factors]
     merge_mertens = cv2.createMergeMertens()
     hdr_image = merge_mertens.process(images)
@@ -132,6 +142,8 @@ lazy_pipe = LazyLoadPipeline()
 
 @timer_func
 def gradio_process_image(input_image, resolution, num_inference_steps, strength, hdr, guidance_scale):
+    if input_image is None:
+        raise gr.Error("Please upload an input image.")
     print("Starting image processing...")
     condition_image = resize_and_upscale(input_image, resolution)
     condition_image = create_hdr_effect(condition_image, hdr)
@@ -160,7 +172,7 @@ def gradio_process_image(input_image, resolution, num_inference_steps, strength,
 
 # Gradio interface
 title = """<h1 align="center">Image Upscaler with Tile ControlNet</h1>
-<p align="center">CPU-optimized version for Hugging Face Spaces</p>"""
+<p align="center">CPU-optimized for Hugging Face Spaces</p>"""
 
 with gr.Blocks() as demo:
     gr.HTML(title)
@@ -171,8 +183,8 @@ with gr.Blocks() as demo:
         with gr.Column():
             output_slider = ImageSlider(label="Before / After", type="numpy")
     with gr.Accordion("Advanced Options", open=False):
-        resolution = gr.Slider(minimum=256, maximum=1024, value=512, step=64, label="Resolution")
-        num_inference_steps = gr.Slider(minimum=1, maximum=20, value=10, step=1, label="Inference Steps")
+        resolution = gr.Slider(minimum=256, maximum=768, value=512, step=64, label="Resolution")
+        num_inference_steps = gr.Slider(minimum=1, maximum=15, value=10, step=1, label="Inference Steps")
         strength = gr.Slider(minimum=0, maximum=1, value=0.4, step=0.01, label="Strength")
         hdr = gr.Slider(minimum=0, maximum=1, value=0, step=0.1, label="HDR Effect")
         guidance_scale = gr.Slider(minimum=0, maximum=10, value=3, step=0.5, label="Guidance Scale")
@@ -181,5 +193,5 @@ with gr.Blocks() as demo:
                      inputs=[input_image, resolution, num_inference_steps, strength, hdr, guidance_scale],
                      outputs=output_slider)
 
-# Launch the app
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()