Merge pull request #128 from priyanshuverma-dev/feat-save-images

feat: save images

app/api_routes.py

@@ -9,7 +9,7 @@
     Response,
 )
 import json, re
-from .models import User, Chatbot, Chat
+from .models import User, Chatbot, Chat, Image
 from sqlalchemy.exc import IntegrityError
 from flask_login import login_user, current_user, login_required
 from typing import Union, List, Optional, Dict
@@ -343,3 +343,19 @@ def api_clear_chats(chatbot_id: int) -> Union[Response, tuple[Response, int]]:
         ),
         200,
     )
+
+
+@api_bp.route("/api/create_image", methods=["POST"])
+@login_required
+def api_create_image() -> Response:
+    """API endpoint to create a new image."""
+    prompt: str = request.form["prompt"]
+
+    image: Image = Image(
+        prompt=prompt,
+        user_id=current_user.uid,
+    )
+
+    db.session.add(image)
+    db.session.commit()
+    return jsonify({"success": True, "message": "Image created."})

app/models.py

@@ -57,3 +57,20 @@ class Chat(db.Model):
 
     def __repr__(self) -> str:
         return f"<Chat: \nQuery: {self.user_query}\nResponse: {self.response}>"
+
+
+class Image(db.Model):
+    __tablename__ = "images"
+
+    id: int = db.Column(db.Integer, primary_key=True)
+    prompt: str = db.Column(db.Text, nullable=False)
+    user_id: int = db.Column(db.Integer, nullable=False)
+    public: bool = db.Column(db.Boolean, default=True)
+
+    def to_dict(self) -> dict:
+        return {
+            "id": self.id,
+            "prompt": self.prompt,
+            "public": self.public,
+            "user_id": self.user_id,
+        }

app/routes.py

@@ -8,9 +8,10 @@
     Response,
 )
 from flask_login import logout_user, current_user, login_required
-from .models import User, Chatbot, Chat
+from .models import User, Chatbot, Chat, Image
 from typing import Union, List
 from .helpers import create_default_chatbots
+from .constants import IMAGE_GEN_API
 
 # Define the blueprint
 bp = Blueprint("routes", __name__)
@@ -149,7 +150,11 @@ def anonymous_chatbot() -> Union[str, Response]:
 def imagine_chatbot() -> Union[str, Response]:
     """Render the chatbot page for the specified chatbot."""
     full_page: bool = request.args.get("full", "true").lower() == "true"
-    return render_template("imagine.html", full_page=full_page)
+    images: List[Image] = Image.query.filter_by(user_id=current_user.id).all()
+    base_url = IMAGE_GEN_API
+    return render_template(
+        "imagine.html", full_page=full_page, images=images, base_url=base_url
+    )
 
 
 @bp.route("/chatbot_hub")

app/static/js/imagine.js

@@ -50,30 +50,31 @@ async function generateImage(event) {
   event.preventDefault(); // Prevent the default form submission
 
   const formData = new FormData(messageForm); // Collect form data
-  const promptText = formData.get("prompt");
-  const imageSrc = `https://image.pollinations.ai/prompt/${promptText}`;
-
-  // Create the new image element dynamically with download button and prompt text
-  const newImageHTML = `
-    <div class="flex justify-start items-center space-x-2 mb-2">
-      <div class="max-w-md bg-white dark:bg-dark dark:text-dark/90 text-gray-900 rounded-xl p-4 drop-shadow-md shadow border border-gray-100 dark:border-darker flex flex-col">
-        <img class="rounded-md" src="${imageSrc}" alt="${promptText}">
-        <p class="text-center mt-2">${promptText}</p>
-        <div class="flex justify-between mt-2">
-          <button type="button" class="bg-blue-500 hover:bg-blue-600 text-white rounded-full p-2 drop-shadow transition duration-200 flex items-center justify-center download-btn" title="Download" data-url="${imageSrc}">
-            <i class="fa-solid fa-download"></i>
-          </button>
-        </div>
-      </div>
-    </div>`;
-
-  // Append the new image to the container
-  document
-    .getElementById("imagesContainer")
-    .insertAdjacentHTML("beforeend", newImageHTML);
-
-  scrollToBottom();
-  // TODO
+
+  try {
+    const response = await fetch(`/api/create_image`, {
+      method: "POST",
+      body: formData,
+    });
+
+    const data = await response.json(); // Assuming the server responds with JSON
+
+    if (!response.ok) {
+      throw new Error(data.message);
+    }
+
+    if (data.success) {
+      loadContent("/imagine");
+      setTimeout(scrollToBottom, 100);
+    } else {
+      // Handle errors returned by the server
+      document.getElementById("error").textContent =
+        data.message || "An error occurred.";
+    }
+  } catch (error) {
+    console.error("failed:", error);
+    document.getElementById("error").textContent = error.message;
+  }
 }
 
 messageForm.addEventListener("submit", generateImage);

app/templates/imagine.html

@@ -26,10 +26,30 @@ <h1 class="text-4xl font-extrabold dark:text-dark text-center">Image Generation<
 
 
   <div id="imagesContainer" class="flex-1 overflow-y-auto p-6 space-y-6 h-full no-scrollbar">
+    {% if images %}
+    {% for image in images %}
+    <div class="flex justify-start items-center space-x-2 mb-2">
+      <div
+        class="max-w-md bg-white dark:bg-dark dark:text-dark/90 text-gray-900 rounded-xl p-4 drop-shadow-md shadow border border-gray-100 dark:border-darker flex flex-col">
+        <img class="rounded-md" src="{{ base_url }}/{{ image.prompt }}" alt="{{ image.prompt }}">
+        <p class="text-center mt-2">{{ image.prompt }}</p>
+        <div class="flex justify-between mt-2">
+          <button type="button"
+            class="bg-blue-500 hover:bg-blue-600 text-white rounded-full p-2 drop-shadow transition duration-200 flex items-center justify-center download-btn"
+            title="Download" data-url="{{ base_url }}/{{ image.prompt }}">
+            <i class="fa-solid fa-download"></i>
+          </button>
+        </div>
+      </div>
+    </div>
 
+    {% endfor %}
+    {% else %}
+    <p class="text-center text-gray-500">No Images generated yet.</p>
+    {% endif %}
 
   </div>
-
+  <div id="error" class="text-red-500"></div>
   <form id="prompt-form" class="flex w-full space-x-3 px-3 py-1 border-t dark:border-darker border-lighter">
     <input type="text" name="prompt" id="prompt" required class="block w-full border rounded-3xl p-3 text-lg
         bg-white dark:bg-dark focus:outline-none dark:border-darker