nada

app.py

@@ -1,7 +1,6 @@
-import spaces
+import os
 import torch
 import imageio
-import os
 import gradio as gr
 from subprocess import getoutput
 from diffusers.schedulers import EulerAncestralDiscreteScheduler
@@ -9,53 +8,60 @@ from transformers import T5EncoderModel, T5Tokenizer
 from allegro.pipelines.pipeline_allegro import AllegroPipeline
 from allegro.models.vae.vae_allegro import AllegroAutoencoderKL3D
 from allegro.models.transformers.transformer_3d_allegro import AllegroTransformer3DModel
-
 from huggingface_hub import snapshot_download
 import spaces  # Import ZeroGPU compatibility module
 
+# Ensure the weights directory exists
 weights_dir = './allegro_weights'
 os.makedirs(weights_dir, exist_ok=True)
 
-is_shared_ui = True if "rhymes-ai-Allegro-textToVideo" in os.environ.get('SPACE_ID', "") else False
+# Check if running in a shared UI environment
+is_shared_ui = "rhymes-ai-Allegro-textToVideo" in os.environ.get('SPACE_ID', "")
 is_gpu_associated = torch.cuda.is_available()
 
+# Download the necessary model files if not in shared UI
 if not is_shared_ui:
     snapshot_download(
         repo_id='rhymes-ai/Allegro',
-        allow_patterns=[
-            'scheduler/**',
-            'text_encoder/**',
-            'tokenizer/**',
-            'transformer/**',
-            'vae/**',
-        ],
         local_dir=weights_dir,
+        allow_patterns=['**']  # Download all required files
     )
 
 if is_gpu_associated:
     gpu_info = getoutput('nvidia-smi')
+    print(f"GPU Info: {gpu_info}")
+
+# Check directory structure
+required_dirs = ['vae', 'text_encoder', 'tokenizer', 'scheduler', 'transformer']
+missing_dirs = [d for d in required_dirs if not os.path.exists(os.path.join(weights_dir, d))]
+if missing_dirs:
+    raise FileNotFoundError(f"The following directories are missing in '{weights_dir}': {missing_dirs}")
 
-#@spaces.GPU(duration=120)  # Request GPU for the entire process
-@spaces.GPU()  # Request GPU for the entire process
+@spaces.GPU(duration=120)  # Request GPU for the entire process
 def process_pipeline(user_prompt, guidance_scale, num_sampling_steps, seed, enable_cpu_offload):
     # Define dtype
     dtype = torch.bfloat16
 
     # Load models
     vae = AllegroAutoencoderKL3D.from_pretrained(
-        "./allegro_weights/vae/", 
+        os.path.join(weights_dir, 'vae/'), 
         torch_dtype=torch.float32
     ).cuda()
     vae.eval()
 
-    text_encoder = T5EncoderModel.from_pretrained("./allegro_weights/text_encoder/", torch_dtype=dtype)
-    text_encoder.eval()
+    text_encoder = T5EncoderModel.from_pretrained(
+        os.path.join(weights_dir, 'text_encoder/'), 
+        torch_dtype=dtype
+    ).eval()
 
-    tokenizer = T5Tokenizer.from_pretrained("./allegro_weights/tokenizer/")
+    tokenizer = T5Tokenizer.from_pretrained(os.path.join(weights_dir, 'tokenizer/'))
 
     scheduler = EulerAncestralDiscreteScheduler()
 
-    transformer = AllegroTransformer3DModel.from_pretrained("./allegro_weights/transformer/", torch_dtype=dtype).cuda()
+    transformer = AllegroTransformer3DModel.from_pretrained(
+        os.path.join(weights_dir, 'transformer/'), 
+        torch_dtype=dtype
+    ).cuda()
     transformer.eval()
 
     allegro_pipeline = AllegroPipeline(
@@ -105,13 +111,11 @@ def process_pipeline(user_prompt, guidance_scale, num_sampling_steps, seed, enab
 
     return save_path
 
-
 # Gradio interface function
 def run_inference(user_prompt, guidance_scale, num_sampling_steps, seed, enable_cpu_offload, progress=gr.Progress(track_tqdm=True)):
     result_path = process_pipeline(user_prompt, guidance_scale, num_sampling_steps, seed, enable_cpu_offload)
     return result_path
 
-
 css = """
 div#col-container {
     margin: 0 auto;