Spaces:
Runtime error
Runtime error
File size: 3,815 Bytes
49d0d1a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 |
import argparse
from transformers import AutoModelForCausalLM, AutoTokenizer
sys_prompt = """
Transform the short prompt into a detailed video-generation caption using this structure:
Opening shot type (long/medium/close-up/extreme close-up/full shot)
Primary subject(s) with vivid attributes (colors, textures, actions, interactions)
Dynamic elements (movement, transitions, or changes over time, e.g., 'gradually lowers,' 'begins to climb,' 'camera moves toward...')
Scene composition (background, environment, spatial relationships)
Lighting/atmosphere (natural/artificial, time of day, mood)
Camera motion (zooms, pans, static/handheld shots) if applicable.
Pattern Summary from Examples:
[Shot Type] of [Subject+Action] + [Detailed Subject Description] + [Environmental Context] + [Lighting Conditions] + [Camera Movement]
One case:
Short prompt: a person is playing football
Long prompt: Medium shot of a young athlete in a red jersey sprinting across a muddy field, dribbling a soccer ball with precise footwork. The player glances toward the goalpost, adjusts their stance, and kicks the ball forcefully into the net. Raindrops fall lightly, creating reflections under stadium floodlights. The camera follows the ball’s trajectory in a smooth pan.
Note: If the subject is stationary, incorporate camera movement to ensure the generated video remains dynamic.
Now expand this short prompt: [{}]. Please only output the final long prompt in English.
"""
import argparse
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
sys_prompt = """...""" # your existing sys_prompt string
class PromptEnhancer:
def __init__(self, model_name="Qwen/Qwen2.5-32B-Instruct"):
# Determine device
if torch.backends.mps.is_available():
device = "mps"
device_map = None # No accelerate offloading, load entirely on MPS
else:
device = "cpu"
device_map = None # Load entirely on CPU
self.device = torch.device(device)
# Load model without accelerate offloading to avoid conflicts
self.model = AutoModelForCausalLM.from_pretrained(
model_name,
torch_dtype=torch.float16 if device == "mps" else torch.float32,
device_map=device_map,
)
# Move model to device manually (since device_map=None)
self.model.to(self.device)
self.tokenizer = AutoTokenizer.from_pretrained(model_name)
def __call__(self, prompt):
prompt = prompt.strip()
prompt = sys_prompt.format(prompt)
messages = [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": prompt}
]
text = self.tokenizer.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=True
)
model_inputs = self.tokenizer([text], return_tensors="pt").to(self.device)
generated_ids = self.model.generate(
**model_inputs,
max_new_tokens=2048,
)
generated_ids = [
output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
]
rewritten_prompt = self.tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
return rewritten_prompt
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument("--prompt", type=str, default="In a still frame, a stop sign")
args = parser.parse_args()
prompt_enhancer = PromptEnhancer()
enhanced_prompt = prompt_enhancer(args.prompt)
print(f'Original prompt: {args.prompt}')
print(f'Enhanced prompt: {enhanced_prompt}')
|