Add support for compatible API services & update ui

app.py

@@ -23,12 +23,16 @@ except ImportError as e:
     print(f"Error importing modules: {e}")
     sys.exit(1)
 
-def set_api_keys(anthropic_key, openai_key):
-    """Securely set API keys in environment"""
+def set_api_keys(anthropic_key, openai_key, anthropic_base_url=None, openai_base_url=None):
+    """Securely set API keys and base URLs in environment"""
     if anthropic_key and anthropic_key.strip():
         os.environ["ANTHROPIC_API_KEY"] = anthropic_key.strip()
     if openai_key and openai_key.strip():
         os.environ["OPENAI_API_KEY"] = openai_key.strip()
+    if anthropic_base_url and anthropic_base_url.strip():
+        os.environ["ANTHROPIC_BASE_URL"] = anthropic_base_url.strip()
+    if openai_base_url and openai_base_url.strip():
+        os.environ["OPENAI_BASE_URL"] = openai_base_url.strip()
 
 AVAILABLE_MODELS = [
     "claude-sonnet-4-20250514",
@@ -44,7 +48,7 @@ def create_job_directory() -> Path:
     job_dir = Path(tempfile.mkdtemp(prefix=f"{dir_name}_"))
     return job_dir
 
-def validate_inputs(pdf_file, logo_file, aff_logo_file, text_model, vision_model, poster_width, poster_height, anthropic_key, openai_key):
+def validate_inputs(pdf_file, logo_file, aff_logo_file, text_model, vision_model, poster_width, poster_height, anthropic_key, openai_key, anthropic_base_url, openai_base_url):
     if not pdf_file:
         return "Please upload PDF paper"
     if not logo_file:
@@ -83,12 +87,12 @@ def validate_inputs(pdf_file, logo_file, aff_logo_file, text_model, vision_model
     
     return None
 
-def generate_poster(pdf_file, logo_file, aff_logo_file, text_model, vision_model, poster_width, poster_height, anthropic_key, openai_key, progress=gr.Progress()):
+def generate_poster(pdf_file, logo_file, aff_logo_file, text_model, vision_model, poster_width, poster_height, anthropic_key, openai_key, anthropic_base_url, openai_base_url, progress=gr.Progress()):
     try:
-        # Set API keys first
-        set_api_keys(anthropic_key, openai_key)
+        # Set API keys and base URLs first
+        set_api_keys(anthropic_key, openai_key, anthropic_base_url, openai_base_url)
         
-        error_msg = validate_inputs(pdf_file, logo_file, aff_logo_file, text_model, vision_model, poster_width, poster_height, anthropic_key, openai_key)
+        error_msg = validate_inputs(pdf_file, logo_file, aff_logo_file, text_model, vision_model, poster_width, poster_height, anthropic_key, openai_key, anthropic_base_url, openai_base_url)
         if error_msg:
             return None, None, f"❌ {error_msg}", ""
         
@@ -193,8 +197,18 @@ with gr.Blocks(title="PosterGen") as demo:
     """)
     
     with gr.Row():
+        # Column 1: Upload Files
         with gr.Column(scale=1):
+            gr.Markdown("### 📄 Upload Files")
+            
+            pdf_file = gr.File(label="PDF Paper", file_types=[".pdf"], type="binary")
             
+            with gr.Row():
+                logo_file = gr.File(label="Conference Logo", file_types=["image"], type="binary")
+                aff_logo_file = gr.File(label="Affiliation Logo", file_types=["image"], type="binary")
+        
+        # Column 2: API Keys, Model Settings, Dimensions, Generate Button
+        with gr.Column(scale=1):
             gr.Markdown("### 🔑 API Keys")
             gr.Markdown("⚠️ Keys are processed securely and not stored")
             
@@ -211,14 +225,18 @@ with gr.Blocks(title="PosterGen") as demo:
                     placeholder="sk-...",
                     info="Required for GPT models"
                 )
-            
-            gr.Markdown("### 📄 Upload Files")
-            
-            pdf_file = gr.File(label="PDF Paper", file_types=[".pdf"], type="binary")
-            
+                
             with gr.Row():
-                logo_file = gr.File(label="Conference Logo", file_types=["image"], type="binary")
-                aff_logo_file = gr.File(label="Affiliation Logo", file_types=["image"], type="binary")
+                anthropic_base_url = gr.Textbox(
+                    label="Anthropic Base URL",
+                    placeholder="https://api.anthropic.com",
+                    info="(Optional) Set the Base URL if using compatible API services"
+                )
+                openai_base_url = gr.Textbox(
+                    label="OpenAI Base URL",
+                    placeholder="https://api.openai.com/v1",
+                    info="(Optional) Set the Base URL if using compatible API services"
+                )
             
             gr.Markdown("### 🤖 Model Settings")
             
@@ -234,8 +252,8 @@ with gr.Blocks(title="PosterGen") as demo:
             
             generate_btn = gr.Button("🚀 Generate Poster", variant="primary", size="lg")
             
+        # Column 3: Status and Download
         with gr.Column(scale=1):
-            
             gr.Markdown("### 📊 Status")
             status_output = gr.Textbox(label="Generation Status", placeholder="Click 'Generate Poster' to start...", lines=6)
             
@@ -248,7 +266,7 @@ with gr.Blocks(title="PosterGen") as demo:
     
     generate_btn.click(
         fn=generate_and_display,
-        inputs=[pdf_file, logo_file, aff_logo_file, text_model, vision_model, poster_width, poster_height, anthropic_key, openai_key],
+        inputs=[pdf_file, logo_file, aff_logo_file, text_model, vision_model, poster_width, poster_height, anthropic_key, openai_key, anthropic_base_url, openai_base_url],
         outputs=[download_file, status_output]
     )
 

utils/langgraph_utils.py

@@ -21,26 +21,41 @@ load_dotenv()
 def create_model(config: ModelConfig):
     """create chat model from config"""
     if config.provider == 'openai':
-        return ChatOpenAI(
-            model_name=config.model_name,
-            temperature=config.temperature,
-            max_tokens=config.max_tokens,
-            api_key=os.getenv('OPENAI_API_KEY')
-        )
+        openai_kwargs = {
+            'model_name': config.model_name,
+            'temperature': config.temperature,
+            'max_tokens': config.max_tokens,
+            'api_key': os.getenv('OPENAI_API_KEY')
+        }
+        base_url = os.getenv('OPENAI_BASE_URL')
+        if base_url:
+            openai_kwargs['base_url'] = base_url
+            
+        return ChatOpenAI(**openai_kwargs)
     elif config.provider == 'anthropic':
-        return ChatAnthropic(
-            model=config.model_name,
-            temperature=config.temperature,
-            max_tokens=config.max_tokens,
-            api_key=os.getenv('ANTHROPIC_API_KEY')
-        )
+        anthropic_kwargs = {
+            'model': config.model_name,
+            'temperature': config.temperature,
+            'max_tokens': config.max_tokens,
+            'api_key': os.getenv('ANTHROPIC_API_KEY')
+        }
+        base_url = os.getenv('ANTHROPIC_BASE_URL')
+        if base_url:
+            anthropic_kwargs['base_url'] = base_url
+            
+        return ChatAnthropic(**anthropic_kwargs)
     elif config.provider == 'google':
-        return ChatGoogleGenerativeAI(
-            model=config.model_name,
-            temperature=config.temperature,
-            max_output_tokens=config.max_tokens,
-            google_api_key=os.getenv('GOOGLE_API_KEY')
-        )
+        google_kwargs = {
+            'model': config.model_name,
+            'temperature': config.temperature,
+            'max_output_tokens': config.max_tokens,
+            'google_api_key': os.getenv('GOOGLE_API_KEY')
+        }
+        base_url = os.getenv('GOOGLE_BASE_URL')
+        if base_url:
+            google_kwargs['base_url'] = base_url
+            
+        return ChatGoogleGenerativeAI(**google_kwargs)
     else:
         raise ValueError(f"unsupported provider: {config.provider}")