import os import string import copy import gradio as gr import PIL.Image import torch from transformers import BitsAndBytesConfig, pipeline import re import time import random DESCRIPTION = "# LLaVA 🌋💪 - Now with Arnold Mode and Bodybuilding Coaching Expertise!" model_id = "llava-hf/llava-1.5-7b-hf" quantization_config = BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_compute_dtype=torch.float16 ) pipe = pipeline("image-to-text", model=model_id, model_kwargs={"quantization_config": quantization_config}) bodybuilding_criteria = { "Muscular Size": "Focus on overall muscle mass and development.", "Symmetry": "Ensure balanced development between left and right sides of the body.", "Proportion": "Maintain aesthetically pleasing ratios between different muscle groups.", "Definition": "Achieve clear separation between muscle groups and visible striations.", "Conditioning": "Minimize body fat to enhance muscle definition and vascularity.", "Posing": "Present physique effectively to highlight strengths and minimize weaknesses.", } bodybuilding_tips = [ "Train each muscle group at least twice a week for optimal growth.", "Focus on compound exercises like squats, deadlifts, and bench presses for overall mass.", "Don't neglect your legs! They're half your physique.", "Proper nutrition is key. Eat clean and maintain a caloric surplus for growth.", "Get enough rest. Muscles grow during recovery, not in the gym.", "Practice your posing regularly. It's not just for shows, it helps mind-muscle connection.", "Stay hydrated. Water is crucial for muscle function and recovery.", ] def extract_response_pairs(text): turns = re.split(r'(USER:|ASSISTANT:)', text)[1:] turns = [turn.strip() for turn in turns if turn.strip()] conv_list = [] for i in range(0, len(turns[1::2]), 2): if i + 1 < len(turns[1::2]): conv_list.append([turns[1::2][i].lstrip(":"), turns[1::2][i + 1].lstrip(":")]) return conv_list def add_text(history, text): history = history + [[text, None]] return history, "" # Clear the input field after submission def arnold_speak(text): arnold_phrases = [ "Come with me if you want to lift!", "I'll be back... after my protein shake.", "Hasta la vista, baby weight!", "Get to da choppa... I mean, da squat rack!", "You lack discipline! But don't worry, I'm here to pump you up!" ] text = text.replace(".", "!") # More enthusiastic punctuation text = text.replace("gym", "iron paradise") text = text.replace("exercise", "pump iron") text = text.replace("workout", "sculpt your physique") # Add bodybuilding advice if random.random() < 0.7: # 70% chance to add bodybuilding advice advice = random.choice(list(bodybuilding_criteria.items())) text += f" Remember, in bodybuilding, {advice[0]} is crucial! {advice[1]}" # Add a bodybuilding tip if random.random() < 0.5: # 50% chance to add a tip tip = random.choice(bodybuilding_tips) text += f" Here's a pro tip: {tip}" # Add random Arnold phrase to the end text += " " + random.choice(arnold_phrases) return text def infer(image, prompt, temperature, length_penalty, repetition_penalty, max_length, min_length, top_p): try: outputs = pipe(images=image, prompt=prompt, generate_kwargs={"temperature": temperature, "length_penalty": length_penalty, "repetition_penalty": repetition_penalty, "max_length": max_length, "min_length": min_length, "top_p": top_p}) inference_output = outputs[0]["generated_text"] return inference_output except Exception as e: print(f"Error during inference: {str(e)}") return f"An error occurred during inference: {str(e)}" def bot(history, text_input, image, temperature, length_penalty, repetition_penalty, max_length, min_length, top_p, arnold_mode): if text_input == "": yield history + [["Please input text", None]] return if image is None: yield history + [["Please input image or wait for image to be uploaded before clicking submit.", None]] return if arnold_mode: system_prompt = """You are Arnold Schwarzenegger, the famous bodybuilder, actor, and former Mr. Olympia. Respond in his iconic style, using his catchphrases and focusing on fitness, bodybuilding, and motivation. Incorporate bodybuilding judging criteria and tips in your responses when relevant.""" else: system_prompt = "You are a helpful AI assistant. Provide clear and concise responses to the user's questions about the image and text input." # Use only the current input for generating the response prompt = f"{system_prompt}\nUSER: \n{text_input}\nASSISTANT:" response = infer(image, prompt, temperature, length_penalty, repetition_penalty, max_length, min_length, top_p) if arnold_mode: response = arnold_speak(response) history.append([text_input, ""]) for i in range(len(response)): history[-1][1] = response[:i+1] time.sleep(0.05) yield history with gr.Blocks() as demo: gr.Markdown(DESCRIPTION) gr.Markdown("""## LLaVA, one of the greatest multimodal chat models is now available in Transformers with 4-bit quantization! ⚡️ See the docs here: https://huggingface.co/docs/transformers/main/en/model_doc/llava.""") chatbot = gr.Chatbot() with gr.Row(): image = gr.Image(type="pil") with gr.Column(): text_input = gr.Textbox(label="Chat Input", lines=3) arnold_mode = gr.Checkbox(label="Arnold Schwarzenegger Mode") with gr.Accordion(label="Advanced settings", open=False): temperature = gr.Slider(label="Temperature", minimum=0.5, maximum=1.0, value=1.0, step=0.1) length_penalty = gr.Slider(label="Length Penalty", minimum=-1.0, maximum=2.0, value=1.0, step=0.2) repetition_penalty = gr.Slider(label="Repetition Penalty", minimum=1.0, maximum=5.0, value=1.5, step=0.5) max_length = gr.Slider(label="Max Length", minimum=1, maximum=500, value=200, step=1) min_length = gr.Slider(label="Minimum Length", minimum=1, maximum=100, value=1, step=1) top_p = gr.Slider(label="Top P", minimum=0.5, maximum=1.0, value=0.9, step=0.1) with gr.Row(): clear_button = gr.Button("Clear") submit_button = gr.Button("Submit", variant="primary") submit_button.click( fn=bot, inputs=[chatbot, text_input, image, temperature, length_penalty, repetition_penalty, max_length, min_length, top_p, arnold_mode], outputs=chatbot ).then( fn=lambda: "", outputs=text_input ) clear_button.click(lambda: ([], None), outputs=[chatbot, image], queue=False) examples = [ ["./examples/bodybuilder.jpeg", "What do you think of this physique?"], ["./examples/gym.jpeg", "How can I improve my workout routine?"] ] gr.Examples(examples=examples, inputs=[image, text_input]) if __name__ == "__main__": demo.queue(max_size=10).launch(debug=True)