Update app.py

app.py

@@ -2,13 +2,15 @@ import os
 import sys
 import torch
 import gradio as gr
-import spaces
 from PIL import Image
 import numpy as np
 from omegaconf import OmegaConf
-import requests
-from tqdm import tqdm
 import subprocess
+from tqdm import tqdm
+import requests
+
+# Assuming spaces is a valid module
+import spaces
 
 def download_file(url, filename):
     response = requests.get(url, stream=True)
@@ -45,38 +47,52 @@ def setup_environment():
 
 setup_environment()
 
-from ccsr.models.ccsr import CCSR
-from ccsr.utils.util import instantiate_from_config
+# Importing from the CCSR folder
+from CCSR.ldm.xformers_state import disable_xformers
+from CCSR.model.q_sampler import SpacedSampler
+from CCSR.model.ccsr_stage1 import ControlLDM
+from CCSR.utils.common import instantiate_from_config, load_state_dict
 
 config = OmegaConf.load("configs/model/ccsr_stage2.yaml")
-model = instantiate_from_config(config.model)
+model = instantiate_from_config(config)
 ckpt = torch.load("weights/real-world_ccsr.ckpt", map_location="cpu")
-model.load_state_dict(ckpt["state_dict"], strict=False)
-model.cuda().eval()
+load_state_dict(model, ckpt, strict=True)
+model.freeze()
+model.to("cuda")
+
+@spaces.GPU  # Decorate the inference function with @spaces.GPU
+@torch.no_grad()
+def process(image, steps, t_max, t_min, color_fix_type):
+    image = Image.open(image).convert("RGB")
+    image = image.resize((256, 256), Image.LANCZOS)
+    image = np.array(image)
+
+    sampler = SpacedSampler(model, var_type="fixed_small")
+    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()
 
-@spaces.GPU
-@torch.inference_mode()
-def infer(image, sr_scale, t_max, t_min, color_fix_type):
-    image = Image.open(image).convert("RGB").resize((256, 256), Image.LANCZOS)
-    image = torch.from_numpy(np.array(image)).float().cuda() / 127.5 - 1
-    image = image.permute(2, 0, 1).unsqueeze(0)
+    model.control_scales = [1.0] * 13
 
-    output = model.super_resolution(
-        image,
-        sr_scale=sr_scale,
-        t_max=t_max,
-        t_min=t_min,
-        color_fix_type=color_fix_type
+    height, width = control.size(-2), control.size(-1)
+    shape = (1, 4, height // 8, width // 8)
+    x_T = torch.randn(shape, device=model.device, dtype=torch.float32)
+
+    samples = sampler.sample_ccsr(
+        steps=steps, t_max=t_max, t_min=t_min, shape=shape, cond_img=control,
+        positive_prompt="", negative_prompt="", x_T=x_T,
+        cfg_scale=1.0, color_fix_type=color_fix_type
     )
 
-    output = ((output.squeeze().permute(1, 2, 0).cpu().numpy() + 1) * 127.5).clip(0, 255).astype(np.uint8)
-    return Image.fromarray(output)
+    x_samples = samples.clamp(0, 1)
+    x_samples = (einops.rearrange(x_samples, "b c h w -> b h w c") * 255).cpu().numpy().clip(0, 255).astype(np.uint8)
+
+    return Image.fromarray(x_samples[0])
 
 interface = gr.Interface(
-    fn=infer,
+    fn=process,
     inputs=[
         gr.Image(type="filepath"),
-        gr.Slider(minimum=1, maximum=8, step=1, value=4),
+        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"),