jonathanjordan21's picture
Change to smaller model
c55bf8f
raw
history blame contribute delete
No virus
2.36 kB
from langchain.llms import HuggingFacePipeline
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, AutoModelForSeq2SeqLM
from components import caption_chain, tag_chain
from components import pexels, utils
import os, gc
import gradio as gr
# model = AutoModelForSeq2SeqLM.from_pretrained("declare-lab/flan-alpaca-gpt4-xl")
# tokenizer = AutoTokenizer.from_pretrained("declare-lab/flan-alpaca-gpt4-xl")
model = AutoModelForSeq2SeqLM.from_pretrained("declare-lab/flan-alpaca-large")
tokenizer = AutoTokenizer.from_pretrained("declare-lab/flan-alpaca-large")
pipe = pipeline(
'text2text-generation',
model=model,
tokenizer= tokenizer,
max_length=120
)
local_llm = HuggingFacePipeline(pipeline=pipe)
llm_chain = caption_chain.chain(llm=local_llm)
sum_llm_chain = tag_chain.chain(llm=local_llm)
pexels_api_key = os.getenv('pexels_api_key')
def pred(product_name, orientation):
if orientation == "Shorts/Reels/TikTok (1080 x 1920)":
orientation = "potrait"
height = 1920
width = 1080
elif orientation == "Youtube Videos (1920 x 1080)":
orientation = "landscape"
height = 1080
width = 1920
else :
orientation = "square"
height = 1080
width = 1080
folder_name, sentences = pexels.generate_videos(product_name, pexels_api_key, orientation, height, width, llm_chain, sum_llm_chain)
gc.collect()
utils.combine_videos(folder_name)
return ["\n".join(sentences), os.path.join(folder_name, "Final_Ad_Video.mp4")]
#{'video':os.path.join(folder_name, "Final_Ad_Video.mp4"),
# 'captions':"\n".join(sentences)}
with gr.Blocks() as demo:
gr.Markdown(
"""
# Ads Generator
Create video ads based on your product name using AI
### Note : the video generation takes about 2-4 minutes
"""
)
dimension = gr.Dropdown(
["Shorts/Reels/TikTok (1080 x 1920)", "Facebook/Youtube Videos (1920 x 1080)", "Square (1080 x 1080)"],
label="Video Dimension", info="Choose dimension"
)
product_name = gr.Textbox(label="product name")
captions = gr.Textbox(label="captions")
video = gr.Video()
btn = gr.Button("Submit")
btn.click(pred, inputs=[product_name, dimension], outputs=[captions,video])
demo.launch()