Working on it

endpoints.py

@@ -1,6 +1,6 @@
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
-from routers import inference, training
+from routers import training
 from huggingface_hub import login
 from config import settings
 import torch

requirements-fastapi.txt

@@ -12,4 +12,5 @@ diffusers==0.10.2
 torch
 scipy
 ftfy
-accelerate
+accelerate
+uuid

routers/donut_evaluate.py

@@ -1,90 +0,0 @@
-from transformers import DonutProcessor, VisionEncoderDecoderModel
-import locale
-
-import re
-import json
-import torch
-from tqdm.auto import tqdm
-import numpy as np
-from donut import JSONParseEvaluator
-from datasets import load_dataset
-from functools import lru_cache
-import os
-import time
-from config import settings
-
-locale.getpreferredencoding = lambda: "UTF-8"
-
-
-@lru_cache(maxsize=1)
-def prepare_model():
-    processor = DonutProcessor.from_pretrained(settings.processor)
-    model = VisionEncoderDecoderModel.from_pretrained(settings.model)
-
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-
-    model.eval()
-    model.to(device)
-
-    dataset = load_dataset(settings.dataset, split="test")
-
-    return processor, model, device, dataset
-
-
-def run_evaluate_donut():
-    worker_pid = os.getpid()
-    print(f"Handling evaluation request with worker PID: {worker_pid}")
-
-    start_time = time.time()
-
-    output_list = []
-    accs = []
-
-    processor, model, device, dataset = prepare_model()
-
-    for idx, sample in tqdm(enumerate(dataset), total=len(dataset)):
-        # prepare encoder inputs
-        pixel_values = processor(sample["image"].convert("RGB"), return_tensors="pt").pixel_values
-        pixel_values = pixel_values.to(device)
-        # prepare decoder inputs
-        task_prompt = "<s_cord-v2>"
-        decoder_input_ids = processor.tokenizer(task_prompt, add_special_tokens=False, return_tensors="pt").input_ids
-        decoder_input_ids = decoder_input_ids.to(device)
-
-        # autoregressively generate sequence
-        outputs = model.generate(
-            pixel_values,
-            decoder_input_ids=decoder_input_ids,
-            max_length=model.decoder.config.max_position_embeddings,
-            early_stopping=True,
-            pad_token_id=processor.tokenizer.pad_token_id,
-            eos_token_id=processor.tokenizer.eos_token_id,
-            use_cache=True,
-            num_beams=1,
-            bad_words_ids=[[processor.tokenizer.unk_token_id]],
-            return_dict_in_generate=True,
-        )
-
-        # turn into JSON
-        seq = processor.batch_decode(outputs.sequences)[0]
-        seq = seq.replace(processor.tokenizer.eos_token, "").replace(processor.tokenizer.pad_token, "")
-        seq = re.sub(r"<.*?>", "", seq, count=1).strip()  # remove first task start token
-        seq = processor.token2json(seq)
-
-        ground_truth = json.loads(sample["ground_truth"])
-        ground_truth = ground_truth["gt_parse"]
-        evaluator = JSONParseEvaluator()
-        score = evaluator.cal_acc(seq, ground_truth)
-
-        accs.append(score)
-        output_list.append(seq)
-
-    end_time = time.time()
-    processing_time = end_time - start_time
-
-    scores = {"accuracies": accs, "mean_accuracy": np.mean(accs)}
-    print(scores, f"length : {len(accs)}")
-    print("Mean accuracy:", np.mean(accs))
-    print(f"Evaluation done, worker PID: {worker_pid}")
-
-    return scores, np.mean(accs), processing_time

routers/donut_inference.py

@@ -1,60 +0,0 @@
-import re
-import time
-import torch
-from transformers import DonutProcessor, VisionEncoderDecoderModel
-from config import settings
-from functools import lru_cache
-import os
-
-
-@lru_cache(maxsize=1)
-def load_model():
-    processor = DonutProcessor.from_pretrained(settings.processor)
-    model = VisionEncoderDecoderModel.from_pretrained(settings.model)
-
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    model.to(device)
-
-    return processor, model, device
-
-
-def process_document_donut(image):
-    worker_pid = os.getpid()
-    print(f"Handling inference request with worker PID: {worker_pid}")
-
-    start_time = time.time()
-
-    processor, model, device = load_model()
-
-    # prepare encoder inputs
-    pixel_values = processor(image, return_tensors="pt").pixel_values
-
-    # prepare decoder inputs
-    task_prompt = "<s_cord-v2>"
-    decoder_input_ids = processor.tokenizer(task_prompt, add_special_tokens=False, return_tensors="pt").input_ids
-
-    # generate answer
-    outputs = model.generate(
-        pixel_values.to(device),
-        decoder_input_ids=decoder_input_ids.to(device),
-        max_length=model.decoder.config.max_position_embeddings,
-        early_stopping=True,
-        pad_token_id=processor.tokenizer.pad_token_id,
-        eos_token_id=processor.tokenizer.eos_token_id,
-        use_cache=True,
-        num_beams=1,
-        bad_words_ids=[[processor.tokenizer.unk_token_id]],
-        return_dict_in_generate=True,
-    )
-
-    # postprocess
-    sequence = processor.batch_decode(outputs.sequences)[0]
-    sequence = sequence.replace(processor.tokenizer.eos_token, "").replace(processor.tokenizer.pad_token, "")
-    sequence = re.sub(r"<.*?>", "", sequence, count=1).strip()  # remove first task start token
-
-    end_time = time.time()
-    processing_time = end_time - start_time
-
-    print(f"Inference done, worker PID: {worker_pid}")
-
-    return processor.token2json(sequence), processing_time

routers/donut_training.py

@@ -1,393 +0,0 @@
-# !pip install -q git+https://github.com/huggingface/transformers.git datasets sentencepiece
-# !pip install -q pytorch-lightning==1.9.5 wandb
-
-from config import settings
-from datasets import load_dataset
-from transformers import VisionEncoderDecoderConfig
-from transformers import DonutProcessor, VisionEncoderDecoderModel
-
-import json
-import random
-from typing import Any, List, Tuple
-
-import torch
-from torch.utils.data import Dataset
-
-from torch.utils.data import DataLoader
-
-import re
-from nltk import edit_distance
-import numpy as np
-import os
-import time
-
-import pytorch_lightning as pl
-from functools import lru_cache
-
-from pytorch_lightning.loggers import WandbLogger
-from pytorch_lightning.callbacks import Callback
-from config import settings
-
-added_tokens = []
-
-dataset_name = settings.dataset
-base_config_name = settings.base_config
-base_processor_name = settings.base_processor
-base_model_name = settings.base_model
-model_name = settings.model
-
-@lru_cache(maxsize=1)
-def prepare_job():
-    print("Preparing job...")
-
-    dataset = load_dataset(dataset_name)
-
-    max_length = 768
-    image_size = [1280, 960]
-
-    # update image_size of the encoder
-    # during pre-training, a larger image size was used
-    config = VisionEncoderDecoderConfig.from_pretrained(base_config_name)
-    config.encoder.image_size = image_size  # (height, width)
-    # update max_length of the decoder (for generation)
-    config.decoder.max_length = max_length
-    # TODO we should actually update max_position_embeddings and interpolate the pre-trained ones:
-    # https://github.com/clovaai/donut/blob/0acc65a85d140852b8d9928565f0f6b2d98dc088/donut/model.py#L602
-
-    processor = DonutProcessor.from_pretrained(base_processor_name)
-    model = VisionEncoderDecoderModel.from_pretrained(base_model_name, config=config)
-
-    return model, processor, dataset, config, image_size, max_length
-
-
-class DonutDataset(Dataset):
-    """
-    DonutDataset which is saved in huggingface datasets format. (see details in https://huggingface.co/docs/datasets)
-    Each row, consists of image path(png/jpg/jpeg) and gt data (json/jsonl/txt),
-    and it will be converted into input_tensor(vectorized image) and input_ids(tokenized string).
-    Args:
-        dataset_name_or_path: name of dataset (available at huggingface.co/datasets) or the path containing image files and metadata.jsonl
-        max_length: the max number of tokens for the target sequences
-        split: whether to load "train", "validation" or "test" split
-        ignore_id: ignore_index for torch.nn.CrossEntropyLoss
-        task_start_token: the special token to be fed to the decoder to conduct the target task
-        prompt_end_token: the special token at the end of the sequences
-        sort_json_key: whether or not to sort the JSON keys
-    """
-
-    def __init__(
-            self,
-            dataset_name_or_path: str,
-            max_length: int,
-            split: str = "train",
-            ignore_id: int = -100,
-            task_start_token: str = "<s>",
-            prompt_end_token: str = None,
-            sort_json_key: bool = True,
-    ):
-        super().__init__()
-
-        model, processor, dataset, config, image_size, p1 = prepare_job()
-
-        self.max_length = max_length
-        self.split = split
-        self.ignore_id = ignore_id
-        self.task_start_token = task_start_token
-        self.prompt_end_token = prompt_end_token if prompt_end_token else task_start_token
-        self.sort_json_key = sort_json_key
-
-        self.dataset = load_dataset(dataset_name_or_path, split=self.split)
-        self.dataset_length = len(self.dataset)
-
-        self.gt_token_sequences = []
-        for sample in self.dataset:
-            ground_truth = json.loads(sample["ground_truth"])
-            if "gt_parses" in ground_truth:  # when multiple ground truths are available, e.g., docvqa
-                assert isinstance(ground_truth["gt_parses"], list)
-                gt_jsons = ground_truth["gt_parses"]
-            else:
-                assert "gt_parse" in ground_truth and isinstance(ground_truth["gt_parse"], dict)
-                gt_jsons = [ground_truth["gt_parse"]]
-
-            self.gt_token_sequences.append(
-                [
-                    self.json2token(
-                        gt_json,
-                        update_special_tokens_for_json_key=self.split == "train",
-                        sort_json_key=self.sort_json_key,
-                    )
-                    + processor.tokenizer.eos_token
-                    for gt_json in gt_jsons  # load json from list of json
-                ]
-            )
-
-        self.add_tokens([self.task_start_token, self.prompt_end_token])
-        self.prompt_end_token_id = processor.tokenizer.convert_tokens_to_ids(self.prompt_end_token)
-
-    def json2token(self, obj: Any, update_special_tokens_for_json_key: bool = True, sort_json_key: bool = True):
-        """
-        Convert an ordered JSON object into a token sequence
-        """
-        if type(obj) == dict:
-            if len(obj) == 1 and "text_sequence" in obj:
-                return obj["text_sequence"]
-            else:
-                output = ""
-                if sort_json_key:
-                    keys = sorted(obj.keys(), reverse=True)
-                else:
-                    keys = obj.keys()
-                for k in keys:
-                    if update_special_tokens_for_json_key:
-                        self.add_tokens([fr"<s_{k}>", fr"</s_{k}>"])
-                    output += (
-                            fr"<s_{k}>"
-                            + self.json2token(obj[k], update_special_tokens_for_json_key, sort_json_key)
-                            + fr"</s_{k}>"
-                    )
-                return output
-        elif type(obj) == list:
-            return r"<sep/>".join(
-                [self.json2token(item, update_special_tokens_for_json_key, sort_json_key) for item in obj]
-            )
-        else:
-            obj = str(obj)
-            if f"<{obj}/>" in added_tokens:
-                obj = f"<{obj}/>"  # for categorical special tokens
-            return obj
-
-    def add_tokens(self, list_of_tokens: List[str]):
-        """
-        Add special tokens to tokenizer and resize the token embeddings of the decoder
-        """
-        model, processor, dataset, config, image_size, p1 = prepare_job()
-
-        newly_added_num = processor.tokenizer.add_tokens(list_of_tokens)
-        if newly_added_num > 0:
-            model.decoder.resize_token_embeddings(len(processor.tokenizer))
-            added_tokens.extend(list_of_tokens)
-
-    def __len__(self) -> int:
-        return self.dataset_length
-
-    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        """
-        Load image from image_path of given dataset_path and convert into input_tensor and labels
-        Convert gt data into input_ids (tokenized string)
-        Returns:
-            input_tensor : preprocessed image
-            input_ids : tokenized gt_data
-            labels : masked labels (model doesn't need to predict prompt and pad token)
-        """
-
-        model, processor, dataset, config, image_size, p1 = prepare_job()
-
-        sample = self.dataset[idx]
-
-        # inputs
-        pixel_values = processor(sample["image"], random_padding=self.split == "train",
-                                 return_tensors="pt").pixel_values
-        pixel_values = pixel_values.squeeze()
-
-        # targets
-        target_sequence = random.choice(self.gt_token_sequences[idx])  # can be more than one, e.g., DocVQA Task 1
-        input_ids = processor.tokenizer(
-            target_sequence,
-            add_special_tokens=False,
-            max_length=self.max_length,
-            padding="max_length",
-            truncation=True,
-            return_tensors="pt",
-        )["input_ids"].squeeze(0)
-
-        labels = input_ids.clone()
-        labels[labels == processor.tokenizer.pad_token_id] = self.ignore_id  # model doesn't need to predict pad token
-        # labels[: torch.nonzero(labels == self.prompt_end_token_id).sum() + 1] = self.ignore_id  # model doesn't need to predict prompt (for VQA)
-        return pixel_values, labels, target_sequence
-
-
-def build_data_loaders():
-    print("Building data loaders...")
-
-    model, processor, dataset, config, image_size, max_length = prepare_job()
-
-    # we update some settings which differ from pretraining; namely the size of the images + no rotation required
-    # source: https://github.com/clovaai/donut/blob/master/config/train_cord.yaml
-    processor.feature_extractor.size = image_size[::-1]  # should be (width, height)
-    processor.feature_extractor.do_align_long_axis = False
-
-    train_dataset = DonutDataset(dataset_name, max_length=max_length,
-                                 split="train", task_start_token="<s_cord-v2>", prompt_end_token="<s_cord-v2>",
-                                 sort_json_key=False,  # cord dataset is preprocessed, so no need for this
-                                 )
-
-    val_dataset = DonutDataset(dataset_name, max_length=max_length,
-                               split="validation", task_start_token="<s_cord-v2>", prompt_end_token="<s_cord-v2>",
-                               sort_json_key=False,  # cord dataset is preprocessed, so no need for this
-                               )
-
-    model.config.pad_token_id = processor.tokenizer.pad_token_id
-    model.config.decoder_start_token_id = processor.tokenizer.convert_tokens_to_ids(['<s_cord-v2>'])[0]
-
-    # feel free to increase the batch size if you have a lot of memory
-    # I'm fine-tuning on Colab and given the large image size, batch size > 1 is not feasible
-    # Set num_workers=4
-    train_dataloader = DataLoader(train_dataset, batch_size=1, shuffle=True, num_workers=4)
-    val_dataloader = DataLoader(val_dataset, batch_size=1, shuffle=False, num_workers=4)
-
-    return train_dataloader, val_dataloader, max_length
-
-
-class DonutModelPLModule(pl.LightningModule):
-    def __init__(self, config, processor, model):
-        super().__init__()
-        self.config = config
-        self.processor = processor
-        self.model = model
-
-        self.train_dataloader, self.val_dataloader, self.max_length = build_data_loaders()
-
-    def training_step(self, batch, batch_idx):
-        pixel_values, labels, _ = batch
-
-        outputs = self.model(pixel_values, labels=labels)
-        loss = outputs.loss
-        self.log_dict({"train_loss": loss}, sync_dist=True)
-        return loss
-
-    def validation_step(self, batch, batch_idx, dataset_idx=0):
-        pixel_values, labels, answers = batch
-        batch_size = pixel_values.shape[0]
-        # we feed the prompt to the model
-        decoder_input_ids = torch.full((batch_size, 1), self.model.config.decoder_start_token_id, device=self.device)
-
-        outputs = self.model.generate(pixel_values,
-                                      decoder_input_ids=decoder_input_ids,
-                                      max_length=self.max_length,
-                                      early_stopping=True,
-                                      pad_token_id=self.processor.tokenizer.pad_token_id,
-                                      eos_token_id=self.processor.tokenizer.eos_token_id,
-                                      use_cache=True,
-                                      num_beams=1,
-                                      bad_words_ids=[[self.processor.tokenizer.unk_token_id]],
-                                      return_dict_in_generate=True, )
-
-        predictions = []
-        for seq in self.processor.tokenizer.batch_decode(outputs.sequences):
-            seq = seq.replace(self.processor.tokenizer.eos_token, "").replace(self.processor.tokenizer.pad_token, "")
-            seq = re.sub(r"<.*?>", "", seq, count=1).strip()  # remove first task start token
-            predictions.append(seq)
-
-        scores = list()
-        for pred, answer in zip(predictions, answers):
-            pred = re.sub(r"(?:(?<=>) | (?=</s_))", "", pred)
-            # NOT NEEDED ANYMORE
-            # answer = re.sub(r"<.*?>", "", answer, count=1)
-            answer = answer.replace(self.processor.tokenizer.eos_token, "")
-            scores.append(edit_distance(pred, answer) / max(len(pred), len(answer)))
-
-            if self.config.get("verbose", False) and len(scores) == 1:
-                print(f"Prediction: {pred}")
-                print(f"    Answer: {answer}")
-                print(f" Normed ED: {scores[0]}")
-
-        return scores
-
-    def validation_epoch_end(self, validation_step_outputs):
-        # I set this to 1 manually
-        # (previously set to len(self.config.dataset_name_or_paths))
-        num_of_loaders = 1
-        if num_of_loaders == 1:
-            validation_step_outputs = [validation_step_outputs]
-        assert len(validation_step_outputs) == num_of_loaders
-        cnt = [0] * num_of_loaders
-        total_metric = [0] * num_of_loaders
-        val_metric = [0] * num_of_loaders
-        for i, results in enumerate(validation_step_outputs):
-            for scores in results:
-                cnt[i] += len(scores)
-                total_metric[i] += np.sum(scores)
-            val_metric[i] = total_metric[i] / cnt[i]
-            val_metric_name = f"val_metric_{i}th_dataset"
-            self.log_dict({val_metric_name: val_metric[i]}, sync_dist=True)
-        self.log_dict({"val_metric": np.sum(total_metric) / np.sum(cnt)}, sync_dist=True)
-
-    def configure_optimizers(self):
-        # TODO add scheduler
-        optimizer = torch.optim.Adam(self.parameters(), lr=self.config.get("lr"))
-
-        return optimizer
-
-    def train_dataloader(self):
-        return self.train_dataloader
-
-    def val_dataloader(self):
-        return self.val_dataloader
-
-
-class PushToHubCallback(Callback):
-    def on_train_epoch_end(self, trainer, pl_module):
-        print(f"Pushing model to the hub, epoch {trainer.current_epoch}")
-        pl_module.model.push_to_hub(model_name,
-                                    commit_message=f"Training in progress, epoch {trainer.current_epoch}")
-
-    def on_train_end(self, trainer, pl_module):
-        print(f"Pushing model to the hub after training")
-        pl_module.processor.push_to_hub(model_name,
-                                        commit_message=f"Training done")
-        pl_module.model.push_to_hub(model_name,
-                                    commit_message=f"Training done")
-
-
-def run_training_donut(max_epochs_param, val_check_interval_param, warmup_steps_param):
-    worker_pid = os.getpid()
-    print(f"Handling training request with worker PID: {worker_pid}")
-
-    start_time = time.time()
-
-    # Set epochs = 30
-    # Set num_training_samples_per_epoch = training set size
-    # Set val_check_interval = 0.4
-    # Set warmup_steps: 425 / 8 = 54, 54 * 10 = 540, 540 * 0.15 = 81
-    config_params = {"max_epochs": max_epochs_param,
-                     "val_check_interval": val_check_interval_param,  # how many times we want to validate during an epoch
-                     "check_val_every_n_epoch": 1,
-                     "gradient_clip_val": 1.0,
-                     "num_training_samples_per_epoch": 425,
-                     "lr": 3e-5,
-                     "train_batch_sizes": [8],
-                     "val_batch_sizes": [1],
-                     # "seed":2022,
-                     "num_nodes": 1,
-                     "warmup_steps": warmup_steps_param,  # 425 / 8 = 54, 54 * 10 = 540, 540 * 0.15 = 81
-                     "result_path": "./result",
-                     "verbose": False,
-                     }
-
-    model, processor, dataset, config, image_size, p1 = prepare_job()
-
-    model_module = DonutModelPLModule(config, processor, model)
-
-    # wandb_logger = WandbLogger(project="sparrow", name="invoices-donut-v5")
-
-    # trainer = pl.Trainer(
-    #     accelerator="gpu",
-    #     devices=1,
-    #     max_epochs=config_params.get("max_epochs"),
-    #     val_check_interval=config_params.get("val_check_interval"),
-    #     check_val_every_n_epoch=config_params.get("check_val_every_n_epoch"),
-    #     gradient_clip_val=config_params.get("gradient_clip_val"),
-    #     precision=16,  # we'll use mixed precision
-    #     num_sanity_val_steps=0,
-    #     # logger=wandb_logger,
-    #     callbacks=[PushToHubCallback()],
-    # )
-
-    # trainer.fit(model_module)
-
-    end_time = time.time()
-    processing_time = end_time - start_time
-
-    print(f"Training done, worker PID: {worker_pid}")
-
-    return processing_time

routers/inference.py

@@ -1,81 +0,0 @@
-from fastapi import APIRouter, File, UploadFile, Form
-from typing import Optional
-from PIL import Image
-import urllib.request
-from io import BytesIO
-from config import settings
-import utils
-import os
-import json
-from routers.donut_inference import process_document_donut
-
-
-router = APIRouter()
-
-def count_values(obj):
-    if isinstance(obj, dict):
-        count = 0
-        for value in obj.values():
-            count += count_values(value)
-        return count
-    elif isinstance(obj, list):
-        count = 0
-        for item in obj:
-            count += count_values(item)
-        return count
-    else:
-        return 1
-
-
-@router.post("/inference")
-async def run_inference(file: Optional[UploadFile] = File(None), image_url: Optional[str] = Form(None),
-                        model_in_use: str = Form('donut'), sparrow_key: str = Form(None)):
-
-    if sparrow_key != settings.sparrow_key:
-        return {"error": "Invalid Sparrow key."}
-
-    result = []
-    if file:
-        # Ensure the uploaded file is a JPG image
-        if file.content_type not in ["image/jpeg", "image/jpg"]:
-            return {"error": "Invalid file type. Only JPG images are allowed."}
-
-        image = Image.open(BytesIO(await file.read()))
-        processing_time = 0
-        if model_in_use == 'donut':
-            result, processing_time = process_document_donut(image)
-        utils.log_stats(settings.inference_stats_file, [processing_time, count_values(result), file.filename, settings.model])
-        print(f"Processing time: {processing_time:.2f} seconds")
-    elif image_url:
-        # test image url: https://raw.githubusercontent.com/katanaml/sparrow/main/sparrow-data/docs/input/invoices/processed/images/invoice_10.jpg
-        with urllib.request.urlopen(image_url) as url:
-            image = Image.open(BytesIO(url.read()))
-
-        processing_time = 0
-        if model_in_use == 'donut':
-            result, processing_time = process_document_donut(image)
-        # parse file name from url
-        file_name = image_url.split("/")[-1]
-        utils.log_stats(settings.inference_stats_file, [processing_time, count_values(result), file_name, settings.model])
-        print(f"Processing time inference: {processing_time:.2f} seconds")
-    else:
-        result = {"info": "No input provided"}
-
-    return result
-
-
-@router.get("/statistics")
-async def get_statistics():
-    file_path = settings.inference_stats_file
-
-    # Check if the file exists, and read its content
-    if os.path.exists(file_path):
-        with open(file_path, 'r') as file:
-            try:
-                content = json.load(file)
-            except json.JSONDecodeError:
-                content = []
-    else:
-        content = []
-
-    return content

routers/training.py

@@ -2,8 +2,6 @@ from fastapi import APIRouter, Form, BackgroundTasks
 from config import settings
 import os
 import json
-from routers.donut_evaluate import run_evaluate_donut
-from routers.donut_training import run_training_donut
 import utils
 import torch
 import requests
@@ -11,6 +9,7 @@ from PIL import Image
 from io import BytesIO
 from pydantic import BaseModel
 import base64
+import uuid
 
 from diffusers import StableDiffusionImg2ImgPipeline
 
@@ -26,25 +25,26 @@ class ActionBody(BaseModel):
     prompt: str
     strength: float
     guidance_scale: float
+    resizeW: int
+    resizeH: int
 
 @router.post("/perform-action")
 async def performAction(actionBody: ActionBody):
     
     response = requests.get(actionBody.url)
     init_image = Image.open(BytesIO(response.content)).convert("RGB")
-    init_image = init_image.resize((768, 512))
+    init_image = init_image.resize((actionBody.resizeW, actionBody.resizeH))
     images = pipe(prompt=actionBody.prompt, image=init_image, strength=actionBody.strength, guidance_scale=actionBody.guidance_scale).images
     print(images)
-    print(images[0])
-
     buffered = BytesIO()
     images[0].save(buffered, format="JPEG")
     img_str = base64.b64encode(buffered.getvalue())
-
-    # images[0].save("fantasy_landscape.png")
+    imgUUID = str(uuid.uuid4())
+    images[0].save(imageUUID+".png")
 
     return {
-        "image": img_str
+        "imageName" : imageUUID+".png",
+        "image": "data:image/jpeg;base64,"+img_str
     }
 
 
@@ -67,5 +67,5 @@ async def hifunction():
     # images[0].save("fantasy_landscape.png")
 
     return {
-        "image": img_str
+        "image": "data:image/jpeg;base64,"+img_str
     }