Update app.py

app.py

@@ -9,6 +9,7 @@ import subprocess
 from tqdm import tqdm
 import requests
 import spaces
+import einops  # Import einops to fix the NameError
 
 def download_file(url, filename):
     response = requests.get(url, stream=True)
@@ -60,7 +61,7 @@ model.to("cuda")
 
 @spaces.GPU
 @torch.no_grad()
-def process(image, steps, t_max, t_min, color_fix_type):
+def process(image, steps, t_max, t_min, color_fix_type, scale):
     image = Image.open(image).convert("RGB")
     image = image.resize((256, 256), Image.LANCZOS)
     image = np.array(image)
@@ -69,7 +70,7 @@ def process(image, steps, t_max, t_min, color_fix_type):
     control = torch.tensor(np.stack([image]) / 255.0, dtype=torch.float32, device=model.device).clamp_(0, 1)
     control = einops.rearrange(control, "n h w c -> n c h w").contiguous()
 
-    model.control_scales = [1.0] * 13
+    model.control_scales = [scale] * 13  # Use the scale parameter
 
     height, width = control.size(-2), control.size(-1)
     shape = (1, 4, height // 8, width // 8)
@@ -89,13 +90,14 @@ def process(image, steps, t_max, t_min, color_fix_type):
 interface = gr.Interface(
     fn=process,
     inputs=[
-        gr.Image(type="filepath"),
-        gr.Slider(minimum=1, maximum=100, step=1, value=45),
-        gr.Slider(minimum=0, maximum=1, step=0.0001, value=0.6667),
-        gr.Slider(minimum=0, maximum=1, step=0.0001, value=0.3333),
-        gr.Dropdown(choices=["adain", "wavelet", "none"], value="adain"),
+        gr.Image(type="filepath", label="Input Image"),
+        gr.Slider(minimum=1, maximum=100, step=1, value=45, label="Steps"),
+        gr.Slider(minimum=0, maximum=1, step=0.0001, value=0.6667, label="T Max"),
+        gr.Slider(minimum=0, maximum=1, step=0.0001, value=0.3333, label="T Min"),
+        gr.Dropdown(choices=["adain", "wavelet", "none"], value="adain", label="Color Fix Type"),
+        gr.Slider(minimum=0, maximum=2, step=0.01, value=1.0, label="Scale"),
     ],
-    outputs=gr.Image(type="pil"),
+    outputs=gr.Image(type="pil", label="Output Image"),
     title="CCSR: Continuous Contrastive Super-Resolution",
 )