Update app.py

app.py

@@ -7,6 +7,12 @@ import random
 from PIL import Image
 import torch
 import asyncio  # Import asyncio
+from skyreelsinfer import TaskType
+from skyreelsinfer.offload import OffloadConfig
+from skyreelsinfer.skyreels_video_infer import SkyReelsVideoInfer
+from huggingface_hub.utils import RepositoryNotFoundError, RevisionNotFoundError, EntryNotFoundError
+from diffusers.utils import export_to_video
+from diffusers.utils import load_image
 
 # os.environ["CUDA_VISIBLE_DEVICES"] = ""  # Uncomment if needed
 os.environ["SAFETENSORS_FAST_GPU"] = "1"
@@ -18,13 +24,7 @@ os.putenv("HF_HUB_ENABLE_HF_TRANSFER", "1")
 predictor_state = gr.State(None)
 device="cuda:0" if torch.cuda.is_available() else "cpu" # Pass device to the constructor
 
-
 def init_predictor(task_type: str):
-    from skyreelsinfer import TaskType
-    from skyreelsinfer.offload import OffloadConfig
-    from skyreelsinfer.skyreels_video_infer import SkyReelsVideoInfer
-    from huggingface_hub.utils import RepositoryNotFoundError, RevisionNotFoundError, EntryNotFoundError
-
     try:
         predictor = SkyReelsVideoInfer(
             task_type=TaskType.I2V if task_type == "i2v" else TaskType.T2V,
@@ -45,22 +45,16 @@ def init_predictor(task_type: str):
         print(f"Error loading model: {e}")
         return None
 
-
 # Make generate_video async
 async def generate_video(prompt, image_file, predictor):
-    from diffusers.utils import export_to_video
-    from diffusers.utils import load_image
-
     if image_file is None:
         return gr.Error("Error: For i2v, provide an image.")
     if not isinstance(prompt, str) or not prompt.strip():
         return gr.Error("Error: Please provide a prompt.")
     if predictor is None:
         return gr.Error("Error: Model not loaded.")
-
     random.seed(time.time())
     seed = int(random.randrange(4294967294))
-
     kwargs = {
         "prompt": prompt,
         "height": 256,
@@ -73,7 +67,6 @@ async def generate_video(prompt, image_file, predictor):
         "negative_prompt": "bad quality, blur",
         "cfg_for": False,
     }
-
     try:
         # Load the image and move it to the correct device *before* inference
         image = load_image(image=image_file.name)
@@ -81,26 +74,21 @@ async def generate_video(prompt, image_file, predictor):
         kwargs["image"] = image
     except Exception as e:
         return gr.Error(f"Image loading error: {e}")
-
     try:
         output = predictor.inference(kwargs)
         frames = output
     except Exception as e:
         return gr.Error(f"Inference error: {e}"), None  # Return None for predictor on error
-
     save_dir = "./result/i2v"  # Consistent directory
     os.makedirs(save_dir, exist_ok=True)
     video_out_file = os.path.join(save_dir, f"{prompt[:100]}_{int(seed)}.mp4")
     print(f"Generating video: {video_out_file}")
-
     try:
         export_to_video(frames, video_out_file, fps=24)
     except Exception as e:
         return gr.Error(f"Video export error: {e}"), None # Return None for predictor
-
     return video_out_file, predictor  # Return updated predictor
 
-
 def display_image(file):
     if file is not None:
         return Image.open(file.name)
@@ -118,20 +106,17 @@ async def main():
         prompt_textbox = gr.Text(label="Prompt")
         generate_button = gr.Button("Generate")
         output_video = gr.Video(label="Output Video")
-
         image_file.change(
             display_image,
             inputs=[image_file],
             outputs=[image_file_preview]
         )
-
         generate_button.click(
             fn=generate_video,
             inputs=[prompt_textbox, image_file, predictor_state],
             outputs=[output_video, predictor_state],  # Output predictor_state
         )
         predictor_state.value = await load_model() # load and set predictor
-
     await demo.launch()
 
 if __name__ == "__main__":