Update app.py

app.py

@@ -6,6 +6,8 @@ import os
 import torch
 import numpy as np
 import yaml
+import cv2
+import tempfile
 from huggingface_hub import hf_hub_download
 #from gradio_imageslider import ImageSlider
 
@@ -55,32 +57,60 @@ print("LMHEAD MODEL CKPT:", LM_MODEL)
 lm_head.load_state_dict(torch.load(LM_MODEL, map_location="cpu"), strict=True)
 
 
-def load_img (filename, norm=True,):
-    img = np.array(Image.open(filename).convert("RGB"))
-    if norm:
-        img = img / 255.
-        img = img.astype(np.float32)
-    return img
-
-
-def process_img (image, prompt):
-    img = np.array(image)
-    img = img / 255.
+def process_frame(frame_bgr, prompt):
+    # BGR → RGB
+    frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
+    img = frame_rgb / 255.0
     img = img.astype(np.float32)
-    y = torch.tensor(img).permute(2,0,1).unsqueeze(0).to(device)
 
-    lm_embd = language_model(prompt)
-    lm_embd = lm_embd.to(device)
+    y = torch.tensor(img).permute(2, 0, 1).unsqueeze(0).to(device)
+
+    lm_embd = language_model(prompt).to(device)
 
     with torch.no_grad():
-        text_embd, deg_pred = lm_head (lm_embd)
+        text_embd, _ = lm_head(lm_embd)
         x_hat = model(y, text_embd)
 
-    restored_img = x_hat.squeeze().permute(1,2,0).clamp_(0, 1).cpu().detach().numpy()
-    restored_img = np.clip(restored_img, 0. , 1.)
+    restored = (
+        x_hat.squeeze()
+        .permute(1, 2, 0)
+        .clamp(0, 1)
+        .cpu()
+        .numpy()
+    )
+
+    restored = (restored * 255).astype(np.uint8)
+    restored_bgr = cv2.cvtColor(restored, cv2.COLOR_RGB2BGR)
+
+    return restored_bgr
+
+
+def process_video(video_path, prompt):
+    cap = cv2.VideoCapture(video_path)
+
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+
+    tmp_out = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
+    out_path = tmp_out.name
+
+    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
+    writer = cv2.VideoWriter(out_path, fourcc, fps, (width, height))
+
+    while True:
+        ret, frame = cap.read()
+        if not ret:
+            break
+
+        restored_frame = process_frame(frame, prompt)
+        writer.write(restored_frame)
+
+    cap.release()
+    writer.release()
+
+    return out_path
 
-    restored_img = (restored_img * 255.0).round().astype(np.uint8)  # float32 to uint8
-    return Image.fromarray(restored_img) #(image, Image.fromarray(restored_img))
 
 
 
@@ -146,16 +176,18 @@ css = """
 """
 
 demo = gr.Interface(
-    fn=process_img,
+    fn=process_video,
     inputs=[
-            gr.Image(type="pil", label="Input", value="images/a4960.jpg"),
-            gr.Text(label="Prompt", value="my colors are too off, make it pop so I can use it in instagram")
+        gr.Video(label="Input Video"),
+        gr.Text(
+            label="Prompt",
+            value="enhance this video and improve visual quality"
+        ),
     ],
-    outputs=[gr.Image(type="pil", label="Ouput")],
+    outputs=gr.Video(label="Output Video"),
     title=title,
     description=description,
     article=article,
-    examples=examples,
     css=css,
 )