Generate video from image using Stable Diffusion in Ruby

video.rb

require 'http'
require 'json'

# Load API Key from ENV
API_KEY = ENV['STABILITY']
raise "API key is missing. Set the STABILITY environment variable." if API_KEY.nil? || API_KEY.empty?

API_URL = 'https://api.stability.ai/v2beta/image-to-video'

# Define the path to your image
IMAGE_PATH = 'oldman_resize.jpeg'  # Update with the path to your image

# Optional parameters
SEED = 42
CFG_SCALE = 2.0  # Ensure this is a float
MOTION_BUCKET_ID = 255  # Adjust motion as needed

# Check if the image file exists
unless File.exist?(IMAGE_PATH)
  raise "File not found: #{IMAGE_PATH}"
end

# Prepare headers (No need to manually specify 'Content-Type')
headers = {
  'Authorization' => "Bearer #{API_KEY}"
}

# Prepare the image as multipart form data
begin
  image_file = File.new(IMAGE_PATH, 'rb')
rescue Errno::ENOENT => e
  raise "Error opening the image file: #{e.message}"
end

# Set up the form data, including optional parameters
form_data = {
  image: HTTP::FormData::File.new(image_file, content_type: 'image/jpeg'),
  seed: SEED,
  cfg_scale: CFG_SCALE,
  motion_bucket_id: MOTION_BUCKET_ID
}

# Make the POST request to start video generation
begin
  response = HTTP.headers(headers).post(API_URL, form: form_data)
rescue HTTP::ConnectionError => e
  raise "Failed to connect to the API: #{e.message}"
rescue HTTP::TimeoutError => e
  raise "Request to the API timed out: #{e.message}"
end

# Check for response status and errors
if response.code != 200
  puts "Error: #{response.code} - #{response.body.to_s}"
  exit
end

# Parse the response to extract the generation ID
begin
  result = JSON.parse(response.body.to_s)
  generation_id = result['id']
  if generation_id.nil?
    raise "Error: No generation ID returned in response."
  else
    puts "Generation ID: #{generation_id}"
  end
rescue JSON::ParserError => e
  raise "Error parsing the response JSON: #{e.message}"
end

# Polling status function to download video in MP4 format
def poll_status(generation_id, api_key, api_url)
  status_url = "#{api_url}/result/#{generation_id}"

  loop do
    begin
      puts "Checking status of Generation ID: #{generation_id}..."
      status_response = HTTP.headers({
        'Authorization' => "Bearer #{api_key}",
        'Accept' => 'video/*'  # Request the video directly
      }).get(status_url)

      case status_response.code
      when 200
        # Video generation is complete, write the MP4 file
        File.open("result.mp4", "wb") do |file|
          file.write(status_response.body)
        end
        puts "Video generation succeeded! Saved as result.mp4"
        return "result.mp4"
      when 202
        # Generation is still in progress
        puts "Generation still in progress... Retrying in 10 seconds."
        sleep 10
      when 404
        # Handle 404 Not Found error
        puts "Error: Generation ID not found. Double-check the ID or the generation request may have failed."
        return nil
      else
        # Handle other error statuses
        puts "Error: #{status_response.code} - #{status_response.body.to_s}"
        return nil
      end
    rescue HTTP::ConnectionError => e
      raise "Failed to connect while polling status: #{e.message}"
    end
  end
end

# Poll the generation status and download the video
begin
  video_file = poll_status(generation_id, API_KEY, API_URL)
  if video_file
    puts "Your video is ready and saved as: #{video_file}"
  else
    puts "There was an issue generating the video."
  end
rescue => e
  puts "Error while polling the status: #{e.message}"
end