Update app.py

app.py

@@ -1,23 +1,24 @@
 import spaces
 import gradio as gr
-import argparse
+import argparse  # Import argparse
 import sys
-import time
 import os
 import random
 import subprocess
-from PIL import Image  # Import PIL.Image
+from PIL import Image  # Keep PIL import
+
+subprocess.run(['sh', './sky.sh'])  # Keep if needed
 
-subprocess.run(['sh', './sky.sh'])  # Keep this if needed for setup
 sys.path.append("./SkyReels-V1")
 
-from skyreelsinfer import TaskType
-from skyreelsinfer.offload import OffloadConfig
-# from skyreelsinfer.skyreels_video_infer import Predictor  # Correct: No Predictor import.
+# Corrected Relative Imports
+from SkyReels-V1.skyreelsinfer import TaskType  # Now imported correctly
+from SkyReels-V1.skyreelsinfer.offload import OffloadConfig
+from SkyReels-V1.skyreelsinfer.skyreels_video_infer import SkyReelsVideoSingleGpuInfer  # Import the class
 from diffusers.utils import export_to_video
-# from diffusers.utils import load_image  # Removed: Use PIL directly
 
 import torch
+import logging
 
 torch.backends.cuda.matmul.allow_tf32 = False
 torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False
@@ -25,204 +26,58 @@ torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False
 torch.backends.cudnn.allow_tf32 = False
 torch.backends.cudnn.deterministic = False
 torch.backends.cudnn.benchmark = False
-# torch.backends.cuda.preferred_blas_library="cublas"
-# torch.backends.cuda.preferred_linalg_library="cusolver"
 torch.set_float32_matmul_precision("highest")
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
-import logging  # Correct: Keep logging
-
-# --- Dummy Classes (Keep these for standalone execution) ---
-class OffloadConfig:
-    def __init__(self, high_cpu_memory=False, parameters_level=False, compiler_transformer=False, compiler_cache=""):
-        self.high_cpu_memory = high_cpu_memory
-        self.parameters_level = parameters_level
-        self.compiler_transformer = compiler_transformer
-        self.compiler_cache = compiler_cache
-
-class TaskType:
-    T2V = 0
-    I2V = 1
-
-class LlamaModel:
-    @staticmethod
-    def from_pretrained(*args, **kwargs):
-        return LlamaModel()
-    def to(self, device):
-        return self
-
-class HunyuanVideoTransformer3DModel:
-    @staticmethod
-    def from_pretrained(*args, **kwargs):
-        return HunyuanVideoTransformer3DModel()
-    def to(self, device):
-        return self
-
-class SkyreelsVideoPipeline:
-    @staticmethod
-    def from_pretrained(*args, **kwargs):
-        return SkyreelsVideoPipeline()
-    def to(self, device):
-        return self
-    def __call__(self, *args, **kwargs):
-        frames = [torch.randn(1, 3, 512, 512)]  # Dummy frames
-        return type('obj', (object,), {'frames' : frames})()
-    class vae:
-        @staticmethod
-        def enable_tiling():
-          return
-
-def quantize_(*args, **kwargs):
-    return
-
-def float8_weight_only():
-    return
-# --- End of Dummy Classes/Functions ---
 logger = logging.getLogger(__name__)
 
-class SkyReelsVideoSingleGpuInfer:  # No more multiprocessing!
-    def __init__(
-        self,
-        task_type: TaskType,
-        model_id: str,
-        quant_model: bool = True,
-        is_offload: bool = True,
-        offload_config: OffloadConfig = OffloadConfig(),
-        enable_cfg_parallel: bool = True,  # Remove world_size, local_rank
-    ):
-        self.task_type = task_type
-        self.model_id = model_id
-        self.quant_model = quant_model
-        self.is_offload = is_offload
-        self.offload_config = offload_config
-        self.enable_cfg_parallel = enable_cfg_parallel  # Keep this
-        self.pipe = None
-        self.is_initialized = False
-        self.gpu_device = None
-
-    def _load_model(self, model_id: str, base_model_id: str = "hunyuanvideo-community/HunyuanVideo", quant_model: bool = True):
-        logger.info(f"load model model_id:{model_id} quan_model:{quant_model}")
-        text_encoder = LlamaModel.from_pretrained(
-            base_model_id, subfolder="text_encoder", torch_dtype=torch.bfloat16
-        ).to("cpu")
-        transformer = HunyuanVideoTransformer3DModel.from_pretrained(
-            model_id, torch_dtype=torch.bfloat16, device="cpu"
-        ).to("cpu")
-
-        if quant_model:
-            quantize_(text_encoder, float8_weight_only())
-            text_encoder.to("cpu")
-            torch.cuda.empty_cache()
-            quantize_(transformer, float8_weight_only())
-            transformer.to("cpu")
-            torch.cuda.empty_cache()
-
-        pipe = SkyreelsVideoPipeline.from_pretrained(
-            base_model_id, transformer=transformer, text_encoder=text_encoder, torch_dtype=torch.bfloat16
-        ).to("cpu")
-        pipe.vae.enable_tiling()
-        torch.cuda.empty_cache()
-        return pipe
-
-    def initialize(self):
-        """Initializes the model and moves it to the GPU."""
-        if self.is_initialized:
-            return
-
-        if not torch.cuda.is_available():
-            raise RuntimeError("CUDA is not available. Cannot initialize model.")
-
-        self.gpu_device = "cuda:0"  # Always cuda:0 in single-GPU case
-        self.pipe = self._load_model(model_id=self.model_id, quant_model=self.quant_model)
-
-        # Simplified: No need for max_batch_dim_size with single GPU
-        if self.is_offload:
-            pass
-        else:
-            self.pipe.to(self.gpu_device)
-
-        if self.offload_config.compiler_transformer:
-            torch._dynamo.config.suppress_errors = True
-            os.environ["TORCHINDUCTOR_FX_GRAPH_CACHE"] = "1"
-            # No world_size in cache directory name
-            os.environ["TORCHINDUCTOR_CACHE_DIR"] = f"{self.offload_config.compiler_cache}"
-            self.pipe.transformer = torch.compile(
-                self.pipe.transformer, mode="max-autotune-no-cudagraphs", dynamic=True
-            )
-            if self.offload_config.compiler_transformer:  # Only warm up if compiling
-                self.warm_up()
-        self.is_initialized = True
+# --- Dummy Classes (Moved to skyreelsinfer/__init__.py) ---
 
-    def warm_up(self):
-      if not self.is_initialized:
-          raise RuntimeError("Model must be initialized before warm-up.")
+# --- Global Variables and Argument Parsing ---
 
-      init_kwargs = {
-            "prompt": "A woman is dancing in a room",
-            "height": 544,
-            "width": 960,
-            "guidance_scale": 6,
-            "num_inference_steps": 1,
-            "negative_prompt": "bad quality",
-            "num_frames": 16,
-            "generator": torch.Generator(self.gpu_device).manual_seed(42),
-            "embedded_guidance_scale": 1.0,
-        }
-      if self.task_type == TaskType.I2V:
-        init_kwargs["image"] = Image.new("RGB", (544,960), color="black") #Dummy
-      self.pipe(**init_kwargs)
-      logger.info("Warm-up complete.")
+_predictor = None
+task_type = TaskType.I2V  # Default task type.  IMPORTANT: Set a default.
 
-    def infer(self, **kwargs):
-        """Handles inference requests."""
-        if not self.is_initialized:
-          self.initialize()
-        if "seed" in kwargs:
-            kwargs["generator"] = torch.Generator(self.gpu_device).manual_seed(kwargs["seed"])
-            del kwargs["seed"]
-        assert (self.task_type == TaskType.I2V and "image" in kwargs) or self.task_type == TaskType.T2V
-        result = self.pipe(**kwargs).frames[0]
-        return result
-
-
-# --- Spaces Integration ---
-
-_predictor = None  # Global variable to hold the predictor
-
-@spaces.GPU(duration=90)  # We DO need @spaces.GPU on init_predictor
+@spaces.GPU(duration=90)
 def init_predictor():
     global _predictor
-    logger = logging.getLogger(__name__)  # Correct: Logger inside function
+    global task_type  # Access the global task_type
+    logger = logging.getLogger(__name__)
 
     if _predictor is None:
+        if task_type == TaskType.I2V:
+            model_id = "Skywork/SkyReels-V1-Hunyuan-I2V"
+        elif task_type == TaskType.T2V:
+            model_id = "your_t2v_model_id"  # Replace with your T2V model ID
+        else:
+            raise ValueError(f"Invalid task_type: {task_type}")
+
         _predictor = SkyReelsVideoSingleGpuInfer(
-            task_type=TaskType.I2V,
-            model_id="Skywork/SkyReels-V1-Hunyuan-I2V",  # Replace!
+            task_type=task_type,
+            model_id=model_id,
             quant_model=True,
             is_offload=True,
             offload_config=OffloadConfig(
                 high_cpu_memory=True,
                 parameters_level=True,
-                compiler_transformer=False,  # Set to True to enable compilation/warm-up
+                compiler_transformer=False,
             ),
         )
-        _predictor.initialize()  # Initialize *after* creation
+        _predictor.initialize()
         logger.info("Predictor initialized")
     else:
         logger.warning("Predictor already initialized (should be rare).")
 
-
-
-@spaces.GPU(duration=90) # Now needed, because we write files.
+@spaces.GPU(duration=90) # Needed, because we are saving a file
 def generate_video(prompt, seed, image=None):
+    global _predictor
     global task_type
-    global _predictor  # Correct: Access global _predictor
 
-    print(f"image:{type(image)}")
     if seed == -1:
-        random.seed(time.time())
+        random.seed()  # Use system time for randomness if seed is -1
         seed = int(random.randrange(4294967294))
+
     kwargs = {
         "prompt": prompt,
         "height": 512,
@@ -235,38 +90,60 @@ def generate_video(prompt, seed, image=None):
         "negative_prompt": "Aerial view, aerial view, overexposed, low quality, deformation, a poor composition, bad hands, bad teeth, bad eyes, bad limbs, distortion",
         "cfg_for": False,
     }
-    assert image is not None, "please input image"
-    # kwargs["image"] = load_image(image=image)  # Removed: load image directly with PIL
-    kwargs["image"] = Image.open(image) # Use PIL.Image.open
+
+    if task_type == TaskType.I2V:
+        assert image is not None, "Please input an image for I2V task."
+        kwargs["image"] = Image.open(image)  # Use PIL.Image.open
+    elif task_type == TaskType.T2V:
+      pass # No image needed.
+    else:
+        raise ValueError("Invalid Tasktype")
 
     if _predictor is None:
         init_predictor()
 
-    output = _predictor.infer(**kwargs)  # Correct: Use _predictor
+    output = _predictor.infer(**kwargs)
 
-    save_dir = f"./result/{task_type}"
+    save_dir = f"./result/{task_type.name}"  # Use task_type.name for directory
     os.makedirs(save_dir, exist_ok=True)
-    video_out_file = f"{save_dir}/{prompt[:100].replace('/','')}_{seed}.mp4"
+    video_out_file = f"{save_dir}/{prompt[:100].replace('/', '')}_{seed}.mp4"
     print(f"generate video, local path: {video_out_file}")
     export_to_video(output, video_out_file, fps=24)
-    return video_out_file, kwargs  # Correct: Return filename, kwargs
+    return video_out_file, kwargs # Return the file path
+
 
 def create_gradio_interface():
-        with gr.Blocks() as demo:
-            with gr.Row():
+    with gr.Blocks() as demo:
+        with gr.Row():
+            with gr.Column():
                 image = gr.Image(label="Upload Image", type="filepath")
                 prompt = gr.Textbox(label="Input Prompt")
-                seed = gr.Number(label="Random Seed", value=-1)
-            submit_button = gr.Button("Generate Video")
-            output_video = gr.Video(label="Generated Video")
-            output_params = gr.Textbox(label="Output Parameters")
-            submit_button.click(
-                fn=generate_video,
-                inputs=[prompt, seed, image],
-                outputs=[output_video, output_params],
-            )
-        return demo
+                seed = gr.Number(label="Random Seed", value=-1)  # Default to -1
+            with gr.Column():
+                submit_button = gr.Button("Generate Video")
+                output_video = gr.Video(label="Generated Video")
+                output_params = gr.Textbox(label="Output Parameters")
+
+        submit_button.click(
+            fn=generate_video,
+            inputs=[prompt, seed, image],
+            outputs=[output_video, output_params],
+        )
+    return demo
+
 
 if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--task_type", type=str, default="i2v", choices=["t2v", "i2v"],
+                        help="Task type, 't2v' for text-to-video, 'i2v' for image-to-video.")
+    args = parser.parse_args()
+
+    # Set the global task_type based on command-line arguments
+    if args.task_type == "t2v":
+        task_type = TaskType.T2V
+    elif args.task_type == "i2v":
+        task_type = TaskType.I2V
+    # No else needed, default is already set
+
     demo = create_gradio_interface()
-    demo.queue().launch() # Add queue for async
+    demo.queue().launch() # Add queue