import requests
import time
import os
import re
# Configuration
API_KEY = os.environ.get("RUNPOD_API_KEY")
BASE_URL = "https://api.runpod.ai/v2"
# Endpoint IDs
QWEN_ENDPOINT = "qwen3-32b-awq"
FLUX_ENDPOINT = "black-forest-labs-flux-1-schnell"
WAN_ENDPOINT = "wan-2-5"
def get_headers():
return {
"Authorization": f"Bearer {API_KEY}",
"Content-Type": "application/json",
}
def poll_for_completion(endpoint, job_id, timeout=300):
"""Poll an async job until completion."""
start_time = time.time()
while time.time() - start_time < timeout:
status_response = requests.get(
f"{BASE_URL}/{endpoint}/status/{job_id}",
headers=get_headers(),
)
status = status_response.json()
if status["status"] == "COMPLETED":
return status
elif status["status"] == "FAILED":
raise Exception(f"Job failed: {status}")
else:
print(f" Status: {status['status']}, waiting...")
time.sleep(5)
raise Exception(f"Job timed out after {timeout} seconds")
def enhance_prompt(simple_prompt):
"""Use Qwen3 32B to enhance a simple prompt into a detailed image description."""
print(f"Enhancing prompt: {simple_prompt}")
response = requests.post(
f"{BASE_URL}/{QWEN_ENDPOINT}/openai/v1/chat/completions",
headers=get_headers(),
json={
"model": "Qwen/Qwen3-32B-AWQ",
"messages": [
{
"role": "system",
"content": "You are an expert at writing prompts for AI image generation. "
"Transform the user's simple idea into a detailed, vivid image description. "
"Include details about lighting, style, composition, and atmosphere. "
"Keep the description under 100 words. Output only the enhanced prompt, "
"nothing else. Do not include any thinking or explanation.",
},
{"role": "user", "content": simple_prompt},
],
"max_tokens": 200,
"temperature": 0.7,
},
)
result = response.json()
enhanced = result["choices"][0]["message"]["content"].strip()
# Remove any <think>...</think> blocks (some models include reasoning)
enhanced = re.sub(r"<think>.*?</think>", "", enhanced, flags=re.DOTALL).strip()
# Also handle unclosed <think> tags
enhanced = re.sub(r"<think>.*", "", enhanced, flags=re.DOTALL).strip()
print(f"Enhanced prompt: {enhanced}")
return enhanced
def generate_image(prompt):
"""Use Flux Schnell to generate an image from the prompt."""
print("Generating image with Flux Schnell...")
# Submit async job
response = requests.post(
f"{BASE_URL}/{FLUX_ENDPOINT}/run",
headers=get_headers(),
json={
"input": {
"prompt": prompt,
"width": 1024,
"height": 1024,
"num_inference_steps": 4,
}
},
)
result = response.json()
job_id = result["id"]
print(f" Job submitted: {job_id}")
# Poll for completion
status = poll_for_completion(FLUX_ENDPOINT, job_id)
image_url = status["output"]["image_url"]
print(f" Image URL: {image_url}")
return image_url
def generate_video(image_url, prompt):
"""Use WAN 2.5 to animate the image into a video."""
print("Generating video with WAN 2.5...")
# Submit the job
response = requests.post(
f"{BASE_URL}/{WAN_ENDPOINT}/run",
headers=get_headers(),
json={
"input": {
"image": image_url,
"prompt": prompt,
"duration": 5,
"size": "1280*720",
}
},
)
result = response.json()
job_id = result["id"]
print(f" Job submitted: {job_id}")
# Poll for completion (video takes longer, so increase timeout)
status = poll_for_completion(WAN_ENDPOINT, job_id, timeout=600)
video_url = status["output"]["result"]
print(f" Video URL: {video_url}")
return video_url
def download_file(url, filename):
"""Download a file from a URL."""
print(f"Downloading to {filename}...")
response = requests.get(url)
with open(filename, "wb") as f:
f.write(response.content)
print(f"Saved: {filename}")
def main():
# Your simple prompt
simple_prompt = "a cat wearing sunglasses"
# Step 1: Enhance the prompt
enhanced_prompt = enhance_prompt(simple_prompt)
# Step 2: Generate the image
image_url = generate_image(enhanced_prompt)
# Step 3: Generate the video
video_url = generate_video(image_url, enhanced_prompt)
# Step 4: Download the results
download_file(image_url, "output_image.png")
download_file(video_url, "output_video.mp4")
print("\nPipeline complete!")
print(f"Original prompt: {simple_prompt}")
print(f"Enhanced prompt: {enhanced_prompt}")
print(f"Image: output_image.png")
print(f"Video: output_video.mp4")
if __name__ == "__main__":
if not API_KEY:
print("Error: Set RUNPOD_API_KEY environment variable")
exit(1)
main()
