Update app.py

app.py

@@ -1,27 +1,5 @@
 import zipfile
-def unzip_content():
-    try:
-        # First try using Python's zipfile
-        print("Attempting to unzip content using Python...")
-        with zipfile.ZipFile('./content.zip', 'r') as zip_ref:
-            zip_ref.extractall('.')
-    except Exception as e:
-        print(f"Python unzip failed: {str(e)}")
-        try:
-            # Fallback to system unzip command
-            print("Attempting to unzip content using system command...")
-            subprocess.run(['unzip', '-o', './content.zip'], check=True)
-        except Exception as e:
-            print(f"System unzip failed: {str(e)}")
-            raise Exception("Failed to unzip content using both methods")
-    print("Content successfully unzipped!")
-
-# Try to unzip content at startup
-try:
-    unzip_content()
-except Exception as e:
-    print(f"Warning: Could not unzip content: {str(e)}")
-
+import subprocess
 import gradio as gr
 import numpy as np
 import torch
@@ -35,45 +13,59 @@ from omegaconf import OmegaConf
 from CLIP import clip
 import os
 import sys
+import tempfile
+import io
+from pathlib import Path
+from PIL import Image
+import cv2
+import imageio
 
-#os.chdir('./taming-transformers')
-#from taming.models.vqgan import VQModel
-#os.chdir('..')
+# Add taming transformers path
 taming_path = os.path.join(os.getcwd(), 'taming-transformers')
 sys.path.append(taming_path)
 from taming.models.vqgan import VQModel
 
-
-from PIL import Image
-import cv2
-import imageio
-
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-def create_video(image_folder='./generated', video_name='morphing_video.mp4'):
-    images = sorted([img for img in os.listdir(image_folder) if img.endswith(".png") or img.endswith(".jpg")])
-    if len(images) == 0:
-        print("No images found in the folder.")
-        return None
+def unzip_content():
+    try:
+        print("Attempting to unzip content using Python...")
+        with zipfile.ZipFile('./content.zip', 'r') as zip_ref:
+            zip_ref.extractall('.')
+    except Exception as e:
+        print(f"Python unzip failed: {str(e)}")
+        try:
+            print("Attempting to unzip content using system command...")
+            subprocess.run(['unzip', '-o', './content.zip'], check=True)
+        except Exception as e:
+            print(f"System unzip failed: {str(e)}")
+            raise Exception("Failed to unzip content using both methods")
+    print("Content successfully unzipped!")
 
-    frame = cv2.imread(os.path.join(image_folder, images[0]))
-    height, width, layers = frame.shape
-    video_writer = imageio.get_writer(video_name, fps=10)
+def create_video(images_list, video_name='morphing_video.mp4'):
+    """Create video from a list of image tensors"""
+    if not images_list:
+        print("No images provided.")
+        return None
 
-    for image in images:
-        img_path = os.path.join(image_folder, image)
-        img = imageio.imread(img_path)
-        video_writer.append_data(img)
+    with tempfile.NamedTemporaryFile(suffix='.mp4', delete=False) as temp_video:
+        video_writer = imageio.get_writer(temp_video.name, fps=10)
+        
+        for img_tensor in images_list:
+            # Convert tensor to numpy array
+            img = img_tensor.cpu().numpy().transpose((1, 2, 0))
+            img = (img * 255).astype('uint8')
+            video_writer.append_data(img)
 
-    video_writer.close()
-    return video_name
+        video_writer.close()
+        return temp_video.name
 
-def save_from_tensors(tensor, output_dir, filename):
+def save_from_tensors(tensor):
+    """Process tensor and return the processed version"""
     img = tensor.clone()
     img = img.mul(255).byte()
     img = img.cpu().numpy().transpose((1, 2, 0))
-    os.makedirs(output_dir, exist_ok=True)
-    Image.fromarray(img).save(os.path.join(output_dir, filename))
+    return img
 
 def norm_data(data):
     return (data.clip(-1, 1) + 1) / 2
@@ -216,35 +208,14 @@ def generate_art(include_text, exclude_text, extras_text, num_iterations):
         res_img, res_z = training_loop(params, optimizer, include_enc, exclude_enc, extras_enc,
                                      vqgan_model, clip_model, w1, w2, total_iter=num_iterations)
 
-        # Save results
-        output_dir = "generated"
-        # Create output directory if it doesn't exist
-        os.makedirs(output_dir, exist_ok=True)
-        
-        # Clear any existing files in the output directory
-        for file in os.listdir(output_dir):
-            file_path = os.path.join(output_dir, file)
-            if os.path.isfile(file_path):
-                os.remove(file_path)
-
-        for i, img in enumerate(res_img):
-            save_from_tensors(img, output_dir, f"generated_image_{i:03d}.png")
-
-        # Create video
-        video_path = create_video()
-        
-        # Delete the generated folder and its contents after creating the video
-        import shutil
-        shutil.rmtree(output_dir)
+        # Create video directly from tensors
+        video_path = create_video(res_img)
         
         return video_path
         
     except Exception as e:
-        # If there's an error, ensure the generated folder is cleaned up
-        if os.path.exists("generated"):
-            import shutil
-            shutil.rmtree("generated")
-        raise e  # Re-raise the exception to be handled by the calling function
+        raise e
+
 def gradio_interface(include_text, exclude_text, extras_text, num_iterations):
     try:
         video_path = generate_art(include_text, exclude_text, extras_text, int(num_iterations))
@@ -252,6 +223,12 @@ def gradio_interface(include_text, exclude_text, extras_text, num_iterations):
     except Exception as e:
         return f"An error occurred: {str(e)}"
 
+# Try to unzip content at startup
+try:
+    unzip_content()
+except Exception as e:
+    print(f"Warning: Could not unzip content: {str(e)}")
+
 # Define and launch the Gradio app
 iface = gr.Interface(
     fn=gradio_interface,
@@ -265,20 +242,17 @@ iface = gr.Interface(
         gr.Number(label="Number of Iterations", 
                  value=200, minimum=1, maximum=1000)
     ],
-    outputs=gr.Video(label="Generated Morphing Video"),
+    outputs=gr.Video(label="Generated Morphing Video", format="mp4", autoplay=True),
     title="VQGAN-CLIP Art Generator",
     css="allow",
     allow_flagging="never",
-    #######
     description = """
 Generate artistic videos using VQGAN-CLIP. Enter your prompts separated by commas and adjust the number of iterations. The model will generate a morphing video based on your inputs.
-
 Note: This application requires GPU access. Please either:
 1. Use the Colab notebook available at https://github.com/SanshruthR/VQGAN-CLIP
 2. Clone this space and enable GPU in your personal copy.
 """)
 
-
 if __name__ == "__main__":
     print("Checking GPU availability:", "GPU AVAILABLE" if torch.cuda.is_available() else "NO GPU FOUND")
     iface.launch()