Update cogvideox/ui/controller.py

cogvideox/ui/controller.py

@@ -31,6 +31,7 @@ from safetensors import safe_open
 from ..data.bucket_sampler import ASPECT_RATIO_512, get_closest_ratio
 from ..utils.utils import save_videos_grid
 
+# version check
 gradio_version = pkg_resources.get_distribution("gradio").version
 gradio_version_is_above_4 = int(gradio_version.split(".")[0]) >= 4
 
@@ -43,6 +44,7 @@ css = """
 }
 """
 
+# Scheduler dictionaries
 ddpm_scheduler_dict = {
     "Euler": EulerDiscreteScheduler,
     "Euler A": EulerAncestralDiscreteScheduler,
@@ -57,10 +59,12 @@ flow_scheduler_dict = {
 }
 all_scheduler_dict = {**ddpm_scheduler_dict, **flow_scheduler_dict}
 
+# alias for backward compatibility
+all_cheduler_dict = all_scheduler_dict
+
 
 class Fun_Controller:
     def __init__(self, GPU_memory_mode, scheduler_dict, weight_dtype, config_path=None):
-        # config dirs
         self.basedir = os.getcwd()
         self.config_dir = os.path.join(self.basedir, "config")
         self.diffusion_transformer_dir = os.path.join(self.basedir, "models", "Diffusion_Transformer")
@@ -81,7 +85,6 @@ class Fun_Controller:
         self.refresh_motion_module()
         self.refresh_personalized_model()
 
-        # model placeholders
         self.tokenizer = None
         self.text_encoder = None
         self.vae = None
@@ -192,12 +195,11 @@ class Fun_Controller:
     def get_height_width_from_reference(
         self, base_resolution, start_image, validation_video, control_video
     ):
-        # Build aspect ratios at this resolution
         aspect_ratio_sizes = {
             k: [x / 512 * base_resolution for x in v] for k, v in ASPECT_RATIO_512.items()
         }
         if self.model_type == "Inpaint":
-            if validation_video is not None:
+            if validation_video:
                 vid = cv2.VideoCapture(validation_video)
                 _, frame = vid.read()
                 w, h = Image.fromarray(frame).size
@@ -256,7 +258,7 @@ class Fun_Controller:
         seed_textbox,
         is_api=False,
     ):
-        # implementation omitted for brevity when running locally
+        # local generation logic (omitted)
         pass
 
 
@@ -283,19 +285,15 @@ def post_eas(
     denoise_strength,
     seed_textbox,
 ):
-    # encode files to base64
+    # helper: encode file to base64
     def _encode(path):
         with open(path, "rb") as f:
             return base64.b64encode(f.read()).decode("utf-8")
 
-    if start_image:
-        start_image = _encode(start_image)
-    if end_image:
-        end_image = _encode(end_image)
-    if validation_video:
-        validation_video = _encode(validation_video)
-    if validation_video_mask:
-        validation_video_mask = _encode(validation_video_mask)
+    if start_image:        start_image = _encode(start_image)
+    if end_image:          end_image = _encode(end_image)
+    if validation_video:   validation_video = _encode(validation_video)
+    if validation_video_mask: validation_video_mask = _encode(validation_video_mask)
 
     datas = {
         "base_model_path": base_model_dropdown,
@@ -321,12 +319,9 @@ def post_eas(
     }
 
     session = requests.Session()
-    # propagate EAS token if provided
-    token = os.environ.get("EAS_TOKEN")
-    if token:
-        session.headers.update({"Authorization": token})
+    if os.environ.get("EAS_TOKEN"):
+        session.headers.update({"Authorization": os.environ["EAS_TOKEN"]})
 
-    # build URL (fallback to local Gradio server if EAS_URL is not set)
     eas_env = os.environ.get("EAS_URL")
     if eas_env:
         base_url = eas_env.rstrip("/")
@@ -336,8 +331,8 @@ def post_eas(
         base_url = f"http://{host}:{port}"
     endpoint = f"{base_url}/cogvideox_fun/infer_forward"
 
-    response = session.post(url=endpoint, json=datas, timeout=300)
-    return response.json()
+    resp = session.post(url=endpoint, json=datas, timeout=300)
+    return resp.json()
 
 
 class Fun_Controller_EAS:
@@ -396,7 +391,6 @@ class Fun_Controller_EAS:
         )
 
         if "base64_encoding" not in outputs:
-            # error path
             return (
                 gr.Image(visible=False, value=None),
                 gr.Video(visible=True, value=None),
@@ -408,26 +402,26 @@ class Fun_Controller_EAS:
         prefix = str(idx).zfill(3)
 
         if is_image or length_slider == 1:
-            file_path = os.path.join(self.savedir_sample, f"{prefix}.png")
-            with open(file_path, "wb") as f:
+            path = os.path.join(self.savedir_sample, f"{prefix}.png")
+            with open(path, "wb") as f:
                 f.write(data)
             if gradio_version_is_above_4:
-                return gr.Image(value=file_path, visible=True), gr.Video(value=None, visible=False), "Success"
+                return gr.Image(value=path, visible=True), gr.Video(value=None, visible=False), "Success"
             else:
                 return (
-                    gr.Image.update(value=file_path, visible=True),
+                    gr.Image.update(value=path, visible=True),
                     gr.Video.update(value=None, visible=False),
                     "Success",
                 )
         else:
-            file_path = os.path.join(self.savedir_sample, f"{prefix}.mp4")
-            with open(file_path, "wb") as f:
+            path = os.path.join(self.savedir_sample, f"{prefix}.mp4")
+            with open(path, "wb") as f:
                 f.write(data)
             if gradio_version_is_above_4:
-                return gr.Image(value=None, visible=False), gr.Video(value=file_path, visible=True), "Success"
+                return gr.Image(value=None, visible=False), gr.Video(value=path, visible=True), "Success"
             else:
                 return (
                     gr.Image.update(value=None, visible=False),
-                    gr.Video.update(value=file_path, visible=True),
+                    gr.Video.update(value=path, visible=True),
                     "Success",
                 )