File size: 1,849 Bytes
32a0eda |
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 |
"""
Created By: ishwor subedi
Date: 2024-08-13
"""
import numpy as np
import torch
from PIL import Image
from transformers import BitsAndBytesConfig
from src.services.image_caption.caption import ImageCaption
from src.services.image_generation.image_generate import ImageGenerator
class ImageProcessingPipeline:
def __init__(self):
quantization_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_compute_dtype=torch.float16
)
self.image_caption = ImageCaption(model_id="llava-hf/llava-1.5-7b-hf", quantization_config=quantization_config)
self.image_generator = ImageGenerator()
def generate_image(self, prompt, negative_prompt, style, use_negative_prompt, num_inference_steps,
num_images_per_prompt, seed, width, height, guidance_scale, randomize_seed) -> Image:
image = self.image_generator.generate_image(prompt=prompt, negative_prompt=negative_prompt, style=style,
use_negative_prompt=use_negative_prompt,
num_inference_steps=num_inference_steps,
num_images_per_prompt=num_images_per_prompt, seed=seed, width=width,
height=height, guidance_scale=guidance_scale,
randomize_seed=randomize_seed)
return image
def generate_caption(self, image, prompt, temperature, length_penalty, repetition_penalty, max_length, min_length,
top_p):
caption = self.image_caption.generate([], prompt, image, temperature, length_penalty, repetition_penalty,
max_length, min_length, top_p)
return caption
|