|
|
""" |
|
|
Created By: ishwor subedi |
|
|
Date: 2024-08-13 |
|
|
""" |
|
|
from PIL import Image |
|
|
from fastapi import HTTPException, UploadFile, File |
|
|
from fastapi.routing import APIRouter |
|
|
from six import BytesIO |
|
|
|
|
|
from src.models.models import ImageCaptionRequest, ImageGenerationRequest, LanguageTranslationRequest |
|
|
from src.pipeline.image_processing_pipeline import ImageProcessingPipeline |
|
|
|
|
|
image_processing_pipeline = ImageProcessingPipeline() |
|
|
image_processing_router = APIRouter() |
|
|
|
|
|
|
|
|
@image_processing_router.post("/generate_image") |
|
|
async def generate_image(request: ImageGenerationRequest): |
|
|
try: |
|
|
image = image_processing_pipeline.generate_image( |
|
|
request.prompt, |
|
|
request.negative_prompt, |
|
|
request.style, |
|
|
request.use_negative_prompt, |
|
|
request.num_inference_steps, |
|
|
request.num_images_per_prompt, |
|
|
request.seed, |
|
|
request.width, |
|
|
request.height, |
|
|
request.guidance_scale, |
|
|
request.randomize_seed |
|
|
) |
|
|
base_64_image = BytesIO() |
|
|
image.save(base_64_image, format='PNG') |
|
|
base_64_image = base_64_image.getvalue() |
|
|
json = {'image': base_64_image, 'status_code': 200} |
|
|
|
|
|
return json |
|
|
except Exception as e: |
|
|
raise HTTPException(status_code=500, detail=str(e)) |
|
|
|
|
|
|
|
|
@image_processing_router.post("/generate_caption") |
|
|
async def generate_caption(request: ImageCaptionRequest, image: UploadFile = File(...)): |
|
|
image = Image.open(BytesIO(await image.read())).convert('RGB') |
|
|
|
|
|
try: |
|
|
caption = image_processing_pipeline.generate_caption( |
|
|
image, |
|
|
request.prompt, |
|
|
request.temperature, |
|
|
request.length_penalty, |
|
|
request.repetition_penalty, |
|
|
request.max_length, |
|
|
request.min_length, |
|
|
request.top_p |
|
|
) |
|
|
json = {'caption': caption, 'status_code': 200} |
|
|
return json |
|
|
except Exception as e: |
|
|
raise HTTPException(status_code=500, detail=str(e)) |
|
|
|
|
|
|
