Spaces:
Running
Running
import logging | |
import os | |
from typing import Dict | |
from pathlib import Path | |
from dotenv import load_dotenv | |
from ontology_dc8f06af066e4a7880a5938933236037.config import ConfigClass | |
from ontology_dc8f06af066e4a7880a5938933236037.input import InputClass | |
from ontology_dc8f06af066e4a7880a5938933236037.output import OutputClass | |
from openfabric_pysdk.context import State | |
from core.stub import Stub | |
from core.pipeline import CreativePipeline | |
load_dotenv() | |
# Configure logging | |
logging.basicConfig( | |
level=logging.INFO, | |
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", | |
) | |
logger = logging.getLogger(__name__) | |
# Configurations for the app | |
configurations: Dict[str, ConfigClass] = dict() | |
############################################################ | |
# Config callback function | |
############################################################ | |
def config(configuration: Dict[str, ConfigClass], state: State) -> None: | |
""" | |
Stores user-specific configuration data. | |
Args: | |
configuration (Dict[str, ConfigClass]): A mapping of user IDs to configuration objects. | |
state (State): The current state of the application (not used in this implementation). | |
""" | |
for uid, conf in configuration.items(): | |
logging.info(f"Saving new config for user with id:'{uid}'") | |
configurations[uid] = conf | |
############################################################ | |
# Execution callback function | |
############################################################ | |
def execute(request_data) -> None: | |
""" | |
Main execution entry point for handling a model pass. | |
Args: | |
request_data: The object containing request and response structures. | |
""" | |
# Retrieve input | |
request: InputClass = request_data.request | |
user_prompt = request.prompt | |
# Log the incoming request | |
logger.info(f"Received request with prompt: '{user_prompt}'") | |
# Retrieve user config | |
user_config: ConfigClass = configurations.get("super-user", None) | |
logger.info(f"Using configuration: {configurations}") | |
# Initialize the Stub with app IDs | |
app_ids = user_config.app_ids if user_config else [] | |
# Make sure app IDs are available | |
if not app_ids: | |
text_to_image_app_id = os.environ.get("TEXT_TO_IMAGE_APP_ID") | |
image_to_3d_app_id = os.environ.get("IMAGE_TO_3D_APP_ID") | |
app_ids = [text_to_image_app_id, image_to_3d_app_id] | |
logger.info( | |
f"No app_ids found in config, using environment defaults: {app_ids}" | |
) | |
stub = Stub(app_ids) | |
# Create the creative pipeline | |
pipeline = CreativePipeline(stub) | |
# Execute the creative pipeline | |
try: | |
logger.info(f"Executing creative pipeline for prompt: '{user_prompt}'") | |
result = pipeline.create(prompt=user_prompt) | |
if result.success: | |
response_message = ( | |
f"Created successfully! From your prompt '{user_prompt}', " | |
f"I generated an image and a 3D model." | |
) | |
logger.info(f"Pipeline completed successfully: {result.to_dict()}") | |
else: | |
if result.image_path: | |
response_message = ( | |
f"Partially completed. I was able to generate an image from " | |
f"your prompt '{user_prompt}', but couldn't create the 3D model." | |
) | |
logger.warning(f"Pipeline partially completed: {result.to_dict()}") | |
else: | |
response_message = ( | |
f"I'm sorry, I couldn't process your request '{user_prompt}'. " | |
f"Please try again with a different description." | |
) | |
logger.error(f"Pipeline failed: {result.to_dict()}") | |
except Exception as e: | |
logger.error(f"Error executing pipeline: {e}") | |
response_message = f"An error occurred while processing your request: {str(e)}" | |
# Prepare response | |
response: OutputClass = request_data.response | |
response.message = response_message | |