Docker and和 WebUI

.gitignore

@@ -418,4 +418,5 @@ tags
 [._]*.un~
 .vscode
 .github
-generated_samples/
+generated_samples/
+huggingface

Dockerfile

@@ -0,0 +1,16 @@
+# Use the PyTorch base image
+FROM pytorch/pytorch:latest
+
+# Set the working directory inside the container
+WORKDIR /app
+
+# Copy the current directory into the container
+COPY . /app
+
+# Install necessary Python packages
+RUN pip install -e . && \
+    pip install fastapi python-multipart uvicorn
+
+# Set the entrypoint for the container to launch your FastAPI app
+CMD ["python", "demo/fastapi_app.py"]
+

demo/fastapi_app.py

@@ -1,16 +1,23 @@
 from fastapi import FastAPI, File, Form, UploadFile, HTTPException
 from fastapi.responses import JSONResponse, StreamingResponse
+from fastapi.staticfiles import StaticFiles
 import torch
 from transformers import AutoConfig, AutoModelForCausalLM
 from janus.models import MultiModalityCausalLM, VLChatProcessor
 from PIL import Image
 import numpy as np
 import io
+import os
+
+# Resolve absolute path based on the script's location
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))  # This gets "demo/"
+WEBUI_DIR = os.path.join(BASE_DIR, "webui")  # Moves up to the project root
 
 app = FastAPI()
+app.mount("/webui", StaticFiles(directory=WEBUI_DIR, html=True), name="webui")
 
 # Load model and processor
-model_path = "deepseek-ai/Janus-1.3B"
+model_path = os.getenv("MODEL_NAME", "deepseek-ai/Janus-1.3B")
 config = AutoConfig.from_pretrained(model_path)
 language_config = config.language_config
 language_config._attn_implementation = 'eager'
@@ -63,7 +70,7 @@ def multimodal_understanding(image_data, question, seed, top_p, temperature):
     return answer
 
 
-@app.post("/understand_image_and_question/")
+@app.post("/understand_image_and_question")
 async def understand_image_and_question(
     file: UploadFile = File(...),
     question: str = Form(...),
@@ -152,7 +159,7 @@ def generate_image(prompt, seed, guidance):
         return [Image.fromarray(images[i]).resize((1024, 1024), Image.LANCZOS) for i in range(parallel_size)]
 
 
-@app.post("/generate_images/")
+@app.post("/generate_images")
 async def generate_images(
     prompt: str = Form(...),
     seed: int = Form(None),

demo/webui/APIChat.js

@@ -0,0 +1,164 @@
+document.addEventListener('DOMContentLoaded', function () {
+    const apiUrlInput = document.getElementById('apiUrlInput');
+    const dropdownButton = document.getElementById('dropdownButton');
+    const dropdownMenu = document.getElementById('dropdownMenu');
+    const responsesContainer = document.getElementById('responsesContainer');
+    const promptInput = document.getElementById('promptInput');
+    const sendButton = document.getElementById('sendButton');
+    const errorMessage = document.getElementById('errorMessage');
+
+    const defaultApiUrl = window.location.origin + '/generate_images';
+
+    let apiUrl = localStorage.getItem('lastUsedUrl') || defaultApiUrl;
+    let urlHistory = JSON.parse(localStorage.getItem('urlHistory')) || [];
+
+    // If urlHistory is empty, add the default localhost URL
+    if (urlHistory.length === 0) {
+        saveUrl(defaultApiUrl)
+    }
+
+    apiUrlInput.value = apiUrl;
+    updateDropdown();
+
+    function updateDropdown() {
+        dropdownMenu.innerHTML = urlHistory.length
+            ? urlHistory.map(url => `<div class="dropdown-item">
+                    <span class="select-url">${url}</span>
+                    <button class="delete-button">🗑️</button>
+                </div>`).join('')
+            : '<div class="dropdown-item">No history</div>';
+    }
+
+    function addMessage(sender, body, avatarUrl) {
+        const messageDiv = document.createElement('div');
+        messageDiv.classList.add('chat-message');
+
+        const avatarDiv = document.createElement('div');
+        avatarDiv.classList.add('message-avatar');
+        avatarDiv.innerHTML = avatarUrl ? `<img src="${avatarUrl}" alt="Avatar">` : `<div class="default-avatar">${sender[0]}</div>`;
+
+        const messageContent = document.createElement('div');
+        messageContent.classList.add('message-content');
+
+        // Generate timestamp
+        const timestamp = new Date().toLocaleTimeString([], { hour: '2-digit', minute: '2-digit' });
+
+        messageContent.innerHTML = `
+            <div class="message-header">
+                <span class="message-sender">${sender}</span>
+                <span class="message-timestamp">${timestamp}</span>
+            </div>
+            <div class="message-body"></div>
+        `;
+
+        const messageBody = messageContent.querySelector('.message-body');
+        if (typeof body === 'string') {
+            messageBody.textContent = body;
+        } else {
+            messageBody.appendChild(body);
+        }
+
+        messageDiv.appendChild(avatarDiv);
+        messageDiv.appendChild(messageContent);
+        responsesContainer.appendChild(messageDiv);
+        responsesContainer.scrollTop = responsesContainer.scrollHeight;
+    }
+
+
+    dropdownButton.addEventListener('click', () => {
+        dropdownMenu.style.display = dropdownMenu.style.display === 'none' ? 'block' : 'none';
+    });
+
+    dropdownMenu.addEventListener('click', (e) => {
+        if (e.target.classList.contains('select-url')) {
+            apiUrl = e.target.textContent;
+            apiUrlInput.value = apiUrl;
+            saveUrl(apiUrl);
+            dropdownMenu.style.display = 'none';
+        } else if (e.target.classList.contains('delete-button')) {
+            const parent = e.target.closest('.dropdown-item');
+            const urlToDelete = parent.querySelector('.select-url').textContent;
+            urlHistory = urlHistory.filter(url => url !== urlToDelete);
+            localStorage.setItem('urlHistory', JSON.stringify(urlHistory));
+            updateDropdown();
+        }
+    });
+
+    apiUrlInput.addEventListener('keydown', (e) => {
+        if (e.key === 'Enter') {
+            apiUrl = apiUrlInput.value.trim();  // Ensure the latest value is saved
+            saveUrl(apiUrl);
+        }
+    });
+
+    function saveUrl(url) {
+        if (!url.trim()) return;
+
+        // Set the last used URL
+        localStorage.setItem('lastUsedUrl', url);
+
+        // Add to history if it doesn't exist
+        if (!urlHistory.includes(url)) {
+            urlHistory.push(url);
+            localStorage.setItem('urlHistory', JSON.stringify(urlHistory));
+            updateDropdown();
+        }
+
+        addMessage('System', `API URL set to: ${url}`, './system.png');
+    }
+    async function generateImage() {
+        const prompt = promptInput.value.trim();
+        if (!prompt) return;
+
+        addMessage('You', prompt, './user.png');
+        promptInput.value = '';
+        sendButton.disabled = true;
+        sendButton.textContent = 'Generating...';
+        errorMessage.style.display = 'none';
+
+        try {
+            const params = new URLSearchParams();
+            params.set('prompt', prompt);
+            params.set('seed', Math.floor(Math.random() * 1000000).toString());
+            params.set('guidance', '5');
+
+            const response = await fetch(apiUrl, {
+                method: 'POST',
+                headers: {
+                    'Accept': 'application/json',
+                    'Content-Type': 'application/x-www-form-urlencoded',
+                },
+                body: params.toString(),
+            });
+
+            if (!response.ok) {
+                const errorText = await response.text();
+                throw new Error(`Failed to generate image. Status: ${response.status}, Response: ${errorText}`);
+            }
+
+            const blob = await response.blob();
+            const img = document.createElement('img');
+            img.src = URL.createObjectURL(blob);
+            img.style.maxWidth = '100%';
+            img.style.borderRadius = '8px';
+
+            addMessage('AI', img, './child.jpg');
+        } catch (err) {
+            addMessage('System', `Error generating image: ${err.message}`, './system.png');
+            errorMessage.textContent = `Error generating image: ${err.message}`;
+            errorMessage.style.display = 'block';
+        } finally {
+            sendButton.disabled = false;
+            sendButton.textContent = 'Send';
+        }
+    }
+
+    sendButton.addEventListener('click', generateImage);
+
+    promptInput.addEventListener('keydown', (e) => {
+        if (e.key === 'Enter' && !e.shiftKey) {
+            e.preventDefault();
+            generateImage();
+        }
+    });
+});