Text-to-Any / app.py
KevinSmith94624's picture
Update app.py
219f78d verified
import torch
from diffusers import TextToVideoSDPipeline, DiffusionPipeline
from diffusers.utils import export_to_video
import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
import PIL
from io import BytesIO
from gtts import gTTS
import time
from pydub import AudioSegment
import nltk
from together import Together
import base64
tokenizer = AutoTokenizer.from_pretrained("ParisNeo/LLama-3.2-3B-Lollms-Finetuned-GGUF")
model0 = AutoModelForCausalLM.from_pretrained("ParisNeo/LLama-3.2-3B-Lollms-Finetuned-GGUF", ignore_mismatched_sizes=True)
device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu")
model0 = model0.to(device)
# Initialize Chat History
def chat_with_llama(user_input, chat_history):
# Prepare formatted prompt
prompt = "You are a helpful, respectful and honest general-purpose assistant."
for user_content, assist_content in chat_history:
prompt += f"user: {user_content}\n"
prompt += f"assistant: {assist_content}\n"
prompt += f"user: {user_input}\n'assistant:"
# Tokenize and generate response
inputs = tokenizer(prompt, return_tensors="pt").to("cuda:1")
output = model0.generate(inputs["input_ids"], max_length=4096, max_new_tokens = 1024, temperature=0.7, max_time = 10.0, repetition_penalty = 1.0)
response = tokenizer.decode(output[0], skip_special_tokens=True)
# Extract and append assistant's response
assistant_reply = response.split("assistant:")[-1].split('user:')[0].strip()
chat_history.append((user_input, assistant_reply))
return assistant_reply, chat_history
api_key='YOUR API KEY HERE'
client = Together(api_key=api_key)
def chat_api(user_input, chat_history):
messages = []
for user_content, assist_content in chat_history:
messages += [
{"role":"user", "content":user_content},
{"role":"assistant", "content":assist_content}
]
messages += [{"role":"user", "content":user_input}]
response = client.chat.completions.create(
model="meta-llama/Llama-3.3-70B-Instruct-Turbo",
messages=messages,
)
reply = response.choices[0].message.content
chat_history.append((user_input, reply))
return reply, chat_history
def tti_api(prompt, num_steps = 25, width = 512, heights = 512):
response = client.images.generate(
prompt=prompt,
model="black-forest-labs/FLUX.1-dev",
width=width,
height=heights,
steps=num_steps,
n=1,
response_format="b64_json"
)
image_data = base64.b64decode(response.data[0].b64_json)
return image_data
prompt = 'A nice black lexus 570 car running on the snowy road.'
image = tti_api(prompt, num_steps = 25)
image = PIL.Image.open(BytesIO(image))
image.save('result.png')
image.show()
def ttv(prompt, num_steps = 50):
# Load the text-to-video model from Hugging Face
model_id = "damo-vilab/text-to-video-ms-1.7b" # ModelScope Text-to-Video model
#model_id = "guoyww/animatediff-motion-adapter-v1-5-2" # ModelScope Text-to-Video
pipe = TextToVideoSDPipeline.from_pretrained(model_id, torch_dtype=torch.float16, variant="fp16")
pipe.to("cuda:0") # Use GPU if available
# Generate video frames
print("Generating video... This may take some time.")
with torch.no_grad():
video_frames = pipe(prompt, num_frames=32, height=256, width=256, num_inference_steps=num_steps).frames[0]
# Save the generated video
video_path = export_to_video(video_frames, output_video_path="output_video.mp4")
return video_path
test_video = ttv('An awesome lexus 570 car running on the snowy road, high quality', num_steps = 50)
# Ensure the sentence tokenizer is downloaded (if not already)
nltk.download('punkt')
# Function to convert text to speech and generate SRT content
def tts(text):
# Initialize the Google TTS engine with language (e.g., 'en' for English)
tts = gTTS(text=text, lang='en', slow=False)
# Save to an audio file
audio_path = "output.mp3"
tts.save(audio_path)
# Load the audio file with pydub to get the duration
audio = AudioSegment.from_mp3(audio_path)
duration_ms = len(audio) # Duration in milliseconds
# Split the text into sentences using NLTK
sentences = nltk.sent_tokenize(text)
# Estimate the duration per sentence
chunk_duration_ms = duration_ms // len(sentences) # Estimated duration per sentence
# Generate SRT content
srt_content = ""
start_time = 0 # Start time of the first subtitle
for idx, sentence in enumerate(sentences):
end_time = start_time + chunk_duration_ms
start_time_formatted = time.strftime('%H:%M:%S', time.gmtime(start_time / 1000)) + ',' + f'{start_time % 1000:03d}'
end_time_formatted = time.strftime('%H:%M:%S', time.gmtime(end_time / 1000)) + ',' + f'{end_time % 1000:03d}'
srt_content += f"{idx + 1}\n"
srt_content += f"{start_time_formatted} --> {end_time_formatted}\n"
srt_content += f"{sentence}\n\n"
start_time = end_time # Update start time for the next sentence
return audio_path, srt_content
def tti(prompt, num_steps = 50, width = 512, heights = 512):
# Load the pre-trained Stable Diffusion pipeline from Hugging Face
pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1")
#pipe.load_lora_weights("FradigmaDangerYT/dalle-e-mini")
# Move the pipeline to GPU (you can select the GPU with cuda:1 for the second GPU)
device0 = torch.device("cuda:0") # Use "cuda:0" for the first GPU, "cuda:1" for the second GPU
pipe.to(device0)
print(heights)
# Generate an image
image = pipe(prompt, num_inference_steps = num_steps, width = width, height = heights).images[0] # Generate image from the prompt
return image
prompt = 'A nice black lexus 570 car running on the snowy road.'
image = tti(prompt, num_steps = 25, width = 320, heights = 240)
# image = PIL.Image.open(BytesIO(image))
image.save('result.png')
image.show()
# If demo is on, turn off demo
try:
demo.close()
except:
pass
with gr.Blocks() as demo:
gr.Markdown("""
# Gradio based Text-to-Any Project
""")
with gr.Tab(label="Llama-Chat"):
radios0 = gr.Radio(['use api', 'use loaded model'], value="use api", show_label = False)
gptDialog = gr.Chatbot(label = "Llama-Chat", max_height=512, min_height=512,
autoscroll= True)
with gr.Row(equal_height=True):
prompt0 = gr.Textbox(label = 'Prompt Input', lines = 1, scale = 9, max_lines=2,
autofocus=True, autoscroll=True, placeholder='Type your message here...')
with gr.Column(scale = 1):
generate_btn0 = gr.Button('generate')
clear_btn0 = gr.Button('clear')
with gr.Tab(label="Text-to-Image/Video"):
with gr.Row():
radios1 = gr.Radio(['use api', 'use loaded model'], value="use api", show_label = False)
steps = gr.Slider(value = 50, minimum = 20, maximum = 100, step = 1, label = 'num_steps')
width = gr.Slider(value = 1024, minimum = 240, maximum = 1792, step = 16, label = 'width')
heights = gr.Slider(value = 512, minimum = 160, maximum = 1792, step = 16, label = 'heights')
with gr.Row():
outputImg = gr.Image(type='pil',height= 512, width=512, label="Output Image", interactive=False)
outputVideo = gr.Video(width=512, height=512, label = "Output Video", interactive=False)
with gr.Row(equal_height=True):
prompt1 = gr.Textbox(label = 'Prompt Input', lines = 1, scale = 9, max_lines=2,
autofocus=True, autoscroll=True, placeholder='Type your message here...')
with gr.Column(scale = 1):
generate_btn1 = gr.Button('generate image')
generate_btn11 = gr.Button('generate video')
with gr.Tab(label = "Text-to-Speech"):
outputAudio = gr.Audio(label="Audio Output", interactive = False)
outputSrt = gr.Textbox(label = 'Script Output', lines = 10, max_lines = 5, placeholder = 'Script output here')
with gr.Row(equal_height=False):
prompt2 = gr.Textbox(label = 'Prompt Input', lines = 5, scale = 9, max_lines=5,
autofocus=True, autoscroll=True, placeholder='Type your message here...')
with gr.Column(scale = 1):
generate_btn2 = gr.Button('generate')
clear_btn2 = gr.Button('clear')
with gr.Tab(label = 'About'):
pass
def generate_txt(prompt, check, history):
if check == 'use api':
response, history = chat_api(prompt, history)
if response == None:
gr.Warning('Can not reach api.')
else:
response, history = chat_with_llama(prompt, history)
if response == None:
gr.Warning('Failed to load model.')
return '', history
def clear_chat():
history = []
gr.Info('Cleaned successfully!')
return history
def generate_img(prompt, check, num_steps, width, heights):
if check == 'use api':
image = tti_api(prompt, num_steps = num_steps, width = width, heights = heights)
image = PIL.Image.open(BytesIO(image))
if not image:
gr.Warning('Can not reach api')
gr.Info('Generated Image Successfully!')
else:
image = tti(prompt, num_steps = num_steps, width = width, heights = heights)
gr.Info('Generated Image Successfully!')
return image
def generate_video(prompt, num_steps):
video = ttv(prompt, num_steps)
gr.Info('Generated Video Successfully!')
return video
def generate_speech(prompt):
audio, script = tts(prompt)
gr.Info('Generated Speech Successfully!')
return audio, script
def clear_speech():
gr.Info('Cleaned Successfully!')
return None, ''
prompt0.submit(generate_txt, [prompt0, radios0, gptDialog], [prompt0, gptDialog])
prompt1.submit(generate_img, [prompt1, radios1], [outputImg])
# generate button click event
generate_btn0.click(generate_txt, [prompt0, radios0, gptDialog], [prompt0, gptDialog])
generate_btn1.click(generate_img, [prompt1, radios1, steps, width, heights], [outputImg])
generate_btn11.click(generate_video, [prompt1, steps], [outputVideo])
generate_btn2.click(generate_speech, [prompt2], [outputAudio, outputSrt])
# clear button click event
clear_btn0.click(clear_chat, [], [gptDialog])
clear_btn2.click(clear_speech, [], [outputAudio, outputSrt])
demo.launch(share = True)