Videobook/Videobook.py

import os
import cv2
import numpy as np
import ffmpeg
from math import ceil
from segmindapi import SD2_1, Kadinsky
import gradio as gr
from .TextPreprocessor import TextPreprocessor
from .AVCombiner import AVCombiner
class Videobook:
  def __init__(self):
    self.preprocessor = TextPreprocessor()
    self.combiner = AVCombiner()
  def get_sentences(self, story):
    return self.preprocessor(story)

  def generate_imgs(self, sentences, steps):
    imgs = []
    for sentence in sentences:
      sentence['pos'] = self.style + ' of ' + str(sentence['pos']) + ', ' + self.tags
      imgs.append(self.pipe.generate(prompt = sentence['pos'], negative_prompt = str(sentence['neg']) + " ,nude, disfigured, blurry", num_inference_steps = steps))
    return imgs    

  def generate(self, story, api_key, style, tags, model, steps):
    self.style = style
    self.tags = tags
    if model == "Stable Diffusion v2.1":
      self.pipe = SD2_1(api_key)
    else:
      self.pipe = Kadinsky(api_key)
    processed_sentences, sentences = self.get_sentences(story)
    return AVCombiner()(self.generate_imgs(processed_sentences, steps), sentences, os.getcwd())