Update app.py

app.py

@@ -1,104 +1,88 @@
-import gradio as gr
-from transformers import LLaMAForConditionalGeneration, LLaMATokenizer
+import os
 import torch
-import pandas as pd
-from PyPDF2 import PdfFileReader
-from googleapiclient.discovery import build
-from google_auth_oauthlib.flow import InstalledAppFlow
-from google.auth.transport.requests import Request
-import pickle
-import pydub
+import random
+import math
+import time
+from datetime import datetime
+from typing import Union, List
+from huggingface_hub import hf_hub_download
+import numpy as np
+import PIL.Image
+from diffusers import CogVideoXPipeline, CogVideoXDDIMScheduler, CogVideoXDPMScheduler, VaeImageProcessor
+from diffusers.utils import export_to_video
+import moviepy.editor as mp
 
-# Set up LLaMA model and tokenizer
-model = LLaMAForConditionalGeneration.from_pretrained("facebook/llama-3.1-base")
-tokenizer = LLaMATokenizer.from_pretrained("facebook/llama-3.1-base")
+def download_file(repo_id, filename, subfolder):
+    return hf_hub_download(repo_id=repo_id, filename=filename, subfolder=subfolder)
 
-# Set up Google API credentials
-SCOPES = ['https://www.googleapis.com/auth/drive']
-creds = None
-if creds is None or not creds.valid:
-    if creds and creds.expired and creds.refresh_token:
-        creds.refresh(Request())
-    else:
-        flow = InstalledAppFlow.from_client_secrets_file(
-            'credentials.json', SCOPES)
-        creds = flow.run_local_server(port=0)
-drive_service = build('drive', 'v3', credentials=creds)
+def convert_to_gif(video_path):
+    clip = mp.VideoFileClip(video_path)
+    clip = clip.set_fps(8)
+    clip = clip.resize(height=240)
+    gif_path = video_path.replace(".mp4", ".gif")
+    clip.write_gif(gif_path, fps=8)
+    return gif_path
 
-# Define function to process uploaded files
-def process_file(file):
-    if file.name.endswith('.pdf'):
-        pdf_file = PdfFileReader(file)
-        text = ''
-        for page in range(pdf_file.numPages):
-            text += pdf_file.getPage(page).extractText()
-        return text
-    elif file.name.endswith('.csv') or file.name.endswith('.xlsx'):
-        if file.name.endswith('.csv'):
-            df = pd.read_csv(file)
-        else:
-            df = pd.read_excel(file)
-        return str(df)
-    elif file.name.endswith('.docx'):
-        # You need to implement a function to extract text from Word documents
-        # For simplicity, this example just returns an error message
-        return "Error: Word document support not implemented"
-    elif file.name.endswith('.gsheet'):
-        spreadsheet_id = file.name.split('/')[-1]
-        range_name = 'Sheet1!A1:Z1000'  # You can change this range as needed
-        service = build('sheets', 'v4', credentials=creds)
-        sheet = service.spreadsheets()
-        result = sheet.values().get(spreadsheetId=spreadsheet_id,
-                                    range=range_name).execute()
-        values = result.get('values', [])
-        return str(values)
-    elif file.name.endswith('.gdoc'):
-        document_id = file.name.split('/')[-1]
-        service = build('docs', 'v1', credentials=creds)
-        doc = service.documents().get(documentId=document_id).execute()
-        text = ''
-        for element in doc.get('body').get('content'):
-            if 'paragraph' in element:
-                text += element.get('paragraph').get('elements')[0].get('textRun').get('content')
-        return text
-    elif file.name.endswith('.mp3'):
-        audio = pydub.AudioSegment.from_mp3(file)
-        text = ''
-        # You need to implement a function to transcribe audio
-        # For simplicity, this example just returns an error message
-        return "Error: Audio transcription support not implemented"
-    else:
-        return "Error: File type not supported"
+def save_video(tensor: Union[List[np.ndarray], List[PIL.Image.Image]], fps: int = 8):
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    video_path = f"{timestamp}.mp4"
+    export_to_video(tensor, video_path, fps=fps)
+    return video_path
 
-# Define function to answer questions about the uploaded content
-def answer_question(content, question):
-    inputs = tokenizer.encode(question, return_tensors="pt")
-    outputs = model.generate(inputs, max_length=100)
-    answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return answer
+# Downloading necessary files
+scheduler_config_path = download_file("vdo/CogVideoX-5b", "scheduler_config.json", "scheduler")
+text_encoder_config_path = download_file("vdo/CogVideoX-5b", "config.json", "text_encoder")
+text_encoder_model_1_path = download_file("vdo/CogVideoX-5b", "model-00001-of-00002.safetensors", "text_encoder")
+text_encoder_model_2_path = download_file("vdo/CogVideoX-5b", "model-00002-of-00002.safetensors", "text_encoder")
+text_encoder_index_path = download_file("vdo/CogVideoX-5b", "model.safetensors.index.json", "text_encoder")
+tokenizer_added_tokens_path = download_file("vdo/CogVideoX-5b", "added_tokens.json", "tokenizer")
+tokenizer_special_tokens_map_path = download_file("vdo/CogVideoX-5b", "special_tokens_map.json", "tokenizer")
+tokenizer_model_path = download_file("vdo/CogVideoX-5b", "spiece.model", "tokenizer")
+tokenizer_config_path = download_file("vdo/CogVideoX-5b", "tokenizer_config.json", "tokenizer")
+transformer_config_path = download_file("vdo/CogVideoX-5b", "config.json", "transformer")
+transformer_model_1_path = download_file("vdo/CogVideoX-5b", "diffusion_pytorch_model-00001-of-00002.safetensors", "transformer")
+transformer_model_2_path = download_file("vdo/CogVideoX-5b", "diffusion_pytorch_model-00002-of-00002.safetensors", "transformer")
+transformer_index_path = download_file("vdo/CogVideoX-5b", "diffusion_pytorch_model.safetensors.index.json", "transformer")
+vae_config_path = download_file("vdo/CogVideoX-5b", "config.json", "vae")
+vae_model_path = download_file("vdo/CogVideoX-5b", "diffusion_pytorch_model.safetensors", "vae")
+configuration_path = download_file("vdo/CogVideoX-5b", "configuration.json", "")
+model_index_path = download_file("vdo/CogVideoX-5b", "model_index.json", "")
 
-# Create Gradio interface
-demo = gr.Interface(
-    fn=answer_question,
-    inputs=["file", "text"],
-    outputs="text",
-    title="LLaMA Chatbot",
-    description="Upload a file or paste some text, and ask a question about the content.",
-)
+pipe = CogVideoXPipeline.from_pretrained("/content/CogVideoX-5b", torch_dtype=torch.float16)
+pipe.enable_model_cpu_offload()
+pipe.enable_sequential_cpu_offload()
+pipe.vae.enable_slicing()
+pipe.vae.enable_tiling()
 
-# Define function to update the Gradio interface with the uploaded file's content
-def update_interface(file):
-    content = process_file(file)
-    demo.update(inputs=[content])
+prompt = "A golden retriever, sporting sleek black sunglasses, with its lengthy fur flowing in the breeze, sprints playfully across a rooftop terrace, recently refreshed by a light rain. The scene unfolds from a distance, the dog's energetic bounds growing larger as it approaches the camera, its tail wagging with unrestrained joy, while droplets of water glisten on the concrete behind it. The overcast sky provides a dramatic backdrop, emphasizing the vibrant golden coat of the canine as it dashes towards the viewer."
+seed = 0
 
-# Create Gradio interface with file upload
-demo_with_upload = gr.Interface(
-    fn=update_interface,
-    inputs=["file"],
-    outputs=None,
-    title="LLaMA Chatbot",
-    description="Upload a file to analyze its content.",
-)
+if seed == 0:
+    random.seed(int(time.time()))
+    seed = random.randint(0, 18446744073709551615)
+print(seed)
 
-# Launch Gradio interface
-demo_with_upload.launch()
+with torch.inference_mode():
+    video_pt = pipe(
+        prompt=prompt,
+        num_videos_per_prompt=1,
+        num_inference_steps=50,
+        num_frames=49,
+        use_dynamic_cfg=True,
+        output_type="pt",
+        guidance_scale=7.0,
+        generator=torch.Generator(device="cpu").manual_seed(seed),
+    ).frames
+
+batch_size = video_pt.shape[0]
+batch_video_frames = []
+for batch_idx in range(batch_size):
+    pt_image = video_pt[batch_idx]
+    pt_image = torch.stack([pt_image[i] for i in range(pt_image.shape[0])])
+
+    image_np = VaeImageProcessor.pt_to_numpy(pt_image)
+    image_pil = VaeImageProcessor.numpy_to_pil(image_np)
+    batch_video_frames.append(image_pil)
+
+video_path = save_video(batch_video_frames[0], fps=math.ceil((len(batch_video_frames[0]) - 1) / 6))
+gif_path = convert_to_gif(video_path)