update

model/model_manager.py

@@ -3,7 +3,6 @@ import random
 import gradio as gr
 import requests
 import io, base64, json
-import spaces
 from PIL import Image
 from .models import IMAGE_GENERATION_MODELS, IMAGE_EDITION_MODELS, load_pipeline
 
@@ -21,7 +20,6 @@ class ModelManager:
             pipe = self.loaded_models[model_name]
         return pipe
 
-    @spaces.GPU(duration=60)
     def generate_image_ig(self, prompt, model_name):
         pipe = self.load_model_pipe(model_name)
         result = pipe(prompt=prompt)
@@ -51,15 +49,9 @@ class ModelManager:
                 results.append(result)
         return results[0], results[1]
 
-    @spaces.GPU(duration=150)
     def generate_image_ie(self, textbox_source, textbox_target, textbox_instruct, source_image, model_name):
         pipe = self.load_model_pipe(model_name)
-        if 'PNP' in model_name:
-            result = pipe(src_image = source_image, src_prompt = textbox_source, target_prompt = textbox_target, instruct_prompt = textbox_instruct, num_inversion_steps=100)
-        elif 'Prompt2prompt' in model_name:
-            result = pipe(src_image = source_image, src_prompt = textbox_source, target_prompt = textbox_target, instruct_prompt = textbox_instruct, num_inner_steps=5)
-        else:
-            result = pipe(src_image = source_image, src_prompt = textbox_source, target_prompt = textbox_target, instruct_prompt = textbox_instruct)
+        result = pipe(src_image = source_image, src_prompt = textbox_source, target_prompt = textbox_target, instruct_prompt = textbox_instruct)
         return result
 
     def generate_image_ie_parallel(self, textbox_source, textbox_target, textbox_instruct, source_image, model_A, model_B):

model/models/imagenhub_models.py

@@ -7,5 +7,27 @@ class ImagenHubModel():
     def __call__(self, *args, **kwargs):
         return self.model.infer_one_image(*args, **kwargs)
     
+class PNP(ImagenHubModel):
+    def __init__(self):
+        super().__init__('PNP')
+    
+    def __call__(self, *args, **kwargs):
+        if "num_inversion_steps" not in kwargs:
+            kwargs["num_inversion_steps"] = 100
+        return super().__call__(*args, **kwargs)
+
+class Prompt2prompt(ImagenHubModel):
+    def __init__(self):
+        super().__init__('Prompt2prompt')
+    
+    def __call__(self, *args, **kwargs):
+        if "num_inner_steps" not in kwargs:
+            kwargs["num_inner_steps"] = 5
+        return super().__call__(*args, **kwargs)
+    
 def load_imagenhub_model(model_name, model_type=None):
+    if model_name == 'PNP':
+        return PNP()
+    if model_name == 'Prompt2prompt':
+        return Prompt2prompt()
     return ImagenHubModel(model_name)

serve/constants.py

@@ -13,4 +13,5 @@ LOG_SERVER_ADDR = os.getenv("LOG_SERVER_ADDR", f"{LOG_SERVER}/{LOG_SERVER_SUBDOA
 # LOG SERVER API ENDPOINTS
 APPEND_JSON = "append_json"
 SAVE_IMAGE = "save_image"
+SAVE_LOG = "save_log"
 

serve/log_server.py

@@ -4,9 +4,9 @@ import json
 import os
 import aiofiles
 from .log_utils import build_logger
-from .constants import LOG_SERVER_SUBDOAMIN, APPEND_JSON, SAVE_IMAGE
+from .constants import LOG_SERVER_SUBDOAMIN, APPEND_JSON, SAVE_IMAGE, SAVE_LOG
 
-logger = build_logger("log_server", "log_server.log")
+logger = build_logger("log_server", "log_server.log", add_remote_handler=False)
 
 app = APIRouter(prefix=f"/{LOG_SERVER_SUBDOAMIN}")
 
@@ -37,3 +37,20 @@ async def save_image(image: UploadFile = File(...), image_path: str = Form(...))
         await f.write(content)  # Write the image content to a file
     logger.info(f"Image saved successfully at {image_path}")
     return {"message": f"Image saved successfully at {image_path}"}
+
+@app.post(f"/{SAVE_LOG}")
+async def save_log(message: str = Form(...), log_path: str = Form(...)):
+    """
+    Save a log message to a specified log file on the server.
+    """
+    print(f"Received log message: {message} to be saved at: {log_path}")
+    # Ensure the directory for the log file exists
+    if os.path.dirname(log_path):
+        os.makedirs(os.path.dirname(log_path), exist_ok=True)
+    
+    # Append the log message to the specified log file
+    async with aiofiles.open(log_path, mode='a') as f:
+        await f.write(f"{message}\n")
+    
+    logger.info(f"Romote log message saved to {log_path}")
+    return {"message": f"Log message saved successfully to {log_path}"}

serve/log_utils.py

@@ -10,17 +10,36 @@ import platform
 import sys
 from typing import AsyncGenerator, Generator
 import warnings
+from pathlib import Path
 
 import requests
 
-from .constants import LOGDIR
+from .constants import LOGDIR, LOG_SERVER_ADDR, SAVE_LOG
+from .utils import save_log_str_on_log_server
 
 
 handler = None
 visited_loggers = set()
 
 
-def build_logger(logger_name, logger_filename):
+# Assuming LOGDIR and other necessary imports and global variables are defined
+
+class APIHandler(logging.Handler):
+    """Custom logging handler that sends logs to an API."""
+
+    def __init__(self, apiUrl, log_path, *args, **kwargs):
+        super(APIHandler, self).__init__(*args, **kwargs)
+        self.apiUrl = apiUrl
+        self.log_path = log_path
+
+    def emit(self, record):
+        log_entry = self.format(record)
+        try:
+            save_log_str_on_log_server(log_entry, self.log_path)
+        except requests.RequestException as e:
+            print(f"Error sending log to API: {e}", file=sys.stderr)
+
+def build_logger(logger_name, logger_filename, add_remote_handler=True):
     global handler
 
     formatter = logging.Formatter(
@@ -56,6 +75,15 @@ def build_logger(logger_name, logger_filename):
     # Get logger
     logger = logging.getLogger(logger_name)
     logger.setLevel(logging.INFO)
+    
+    if add_remote_handler:
+        # Add APIHandler to send logs to your API
+        api_url = f"{LOG_SERVER_ADDR}/{SAVE_LOG}"
+        
+        remote_logger_filename = str(Path(logger_filename).stem + "_remote.log")
+        api_handler = APIHandler(apiUrl=api_url, log_path=f"{LOGDIR}/{remote_logger_filename}")
+        api_handler.setFormatter(formatter)
+        logger.addHandler(api_handler)
 
     # if LOGDIR is empty, then don't try output log to local file
     if LOGDIR != "":

serve/utils.py

@@ -6,7 +6,7 @@ import numpy as np
 import gradio as gr
 from pathlib import Path
 from model.model_registry import *
-from .constants import LOGDIR, LOG_SERVER_ADDR, APPEND_JSON, SAVE_IMAGE
+from .constants import LOGDIR, LOG_SERVER_ADDR, APPEND_JSON, SAVE_IMAGE, SAVE_LOG
 from typing import Union
 
 
@@ -159,4 +159,12 @@ def append_json_item_on_log_server(json_item: Union[dict, str], log_file: str):
     url = f"{LOG_SERVER_ADDR}/{APPEND_JSON}"
     # Make the POST request, sending the JSON string and the log file name
     response = requests.post(url, data={'json_str': json_item, 'file_name': log_file})
+    return response
+
+def save_log_str_on_log_server(log_str: str, log_file: str):
+    log_file = Path(log_file).absolute().relative_to(os.getcwd())
+    log_file = str(log_file)
+    url = f"{LOG_SERVER_ADDR}/{SAVE_LOG}"
+    # Make the POST request, sending the log message and the log file name
+    response = requests.post(url, data={'message': log_str, 'log_path': log_file})
     return response