File size: 4,830 Bytes
d01f571
 
9046868
d01f571
0656d52
d01f571
 
 
 
9046868
 
 
 
 
 
 
 
ce97d3a
0656d52
5b02058
0656d52
d01f571
 
 
 
 
0656d52
d01f571
 
 
 
 
0656d52
d01f571
 
 
 
 
 
 
 
0656d52
d01f571
 
 
0656d52
d01f571
 
 
 
 
 
0656d52
d01f571
 
 
 
d745090
4069d7c
d01f571
0656d52
d01f571
0656d52
d01f571
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
677de00
d01f571
 
 
 
7b615e0
 
d01f571
 
 
 
 
 
 
 
583a20f
d01f571
 
 
 
 
 
 
 
 
 
 
 
0656d52
ea87de9
62743c2
9046868
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
from langchain_google_genai import ChatGoogleGenerativeAI
from google import genai
from google.api_core import retry
from io import BytesIO
import gradio as gr
import base64
import time
import uuid
import os

# Define a retry policy
is_retriable = lambda e: (isinstance(e, genai.errors.APIError) and e.code in {429, 503})

if not hasattr(genai.models.Models.generate_content, '__wrapped__'):
  genai.models.Models.generate_content = retry.Retry(
      predicate=is_retriable)(genai.models.Models.generate_content)

os.environ["GOOGLE_API_KEY"] = os.getenv("GOOGLE_API_KEY")

vis_llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash-preview-image-generation")

# Image to Base64 Converter
def image_to_base64(image_path):
  with open(image_path, 'rb') as img:
    encoded_string = base64.b64encode(img.read())
  return encoded_string.decode('utf-8')

# Base64 to Image Converter
def base64_to_image(base64_img, filename):
  with open(filename, "wb") as f:
    f.write(base64.b64decode(base64_img))
  return filename

# Function that takes user inputs and displays it on chatUI
def query_message(message, history):
  if message["files"] is not None:
    for img in message["files"]:
      history.append({"role": "user", "content": {"path": img}})
  if message["text"] is not None:
    history.append({"role": "user", "content": message["text"]})
  return gr.MultimodalTextbox(value=None, interactive=False), history

# Function that takes user inputs, generate response and displays on chatUI
def chat_bot(history):
  messages = history.copy()

  # Format with system prompt
  system_prompt = """
                   system: You are CreativeMuse, a warm, imaginative, and encouraging creative assistant.
                   Your purpose is to help users overcome creative blocks and spark new ideas across various fields
                   — writing, design, music, visual arts, filmmaking, and more. When a user feels stuck, you offer unconventional prompts,
                   inspiring questions, or curated brainstorming techniques to get their creativity flowing again.

                   - Tailor your suggestions to their medium and mood, mixing practical strategies with whimsical, unexpected twists to encourage exploration.
                   Always be supportive, non-judgmental, and affirming—help them feel safe to experiment.
                   You can also suggest small creative exercises, constraints, "what if" scenarios, or draw connections between unrelated concepts to stimulate new thinking.
                   If a user shares part of a project, offer suggestions that expand on their work without taking it over.

                   - Keep your response short and straight to the point. You only generate an image on the user's request. 
                   """

  formatted_prompt = system_prompt + "\n\n"

  # Add conversation history
  for msg in messages:
    if "path" in msg["content"] and isinstance(msg["content"], dict):
      encoded_img = image_to_base64(msg["content"]["path"])
      formatted_prompt += f"{msg['role']}: data:image/jpeg.base64,{encoded_img}"
    else:
      formatted_prompt += f"{msg['role']}: {msg['content']}\n\n"

  # Get response from model
  response = vis_llm.invoke(formatted_prompt,
                            generation_config=dict(response_modalities=["TEXT", "IMAGE"]))

  # Add response to history
  history.append({"role": "assistant", "content": ""})

  for block in response.content:
    if isinstance(block, str):
      history[-1]["content"] += block
      time.sleep(0.02)
      yield history
    elif isinstance(block, dict) and block.get("image_url"):
      url = block["image_url"].get("url").split(",")[-1]
      image = base64_to_image(url, f"{uuid.uuid4()}.png")
      history.append({"role": "assistant", "content": {"path": image}})
      yield history

welcome_message = "Hey, I’m CreativeMuse—your creative companion when inspiration runs dry. Stuck on an idea? Let’s shake things up and spark something new together."

# Interface Code
with gr.Blocks() as app:
  chatbot = gr.Chatbot(
      value=[{"role": "assistant", "content": welcome_message}],
      label="CreativeMuse",
      show_label=True,
      type='messages',
      bubble_full_width=False,
      avatar_images=(None, "1eb05f325ec50a15c8b045f3428d6d5e.jpg")
    )
  text_box = gr.MultimodalTextbox(
          placeholder="Enter text or upload an image...",
          container=False,
          file_types=["image"],
          file_count="multiple",
          interactive=True,
      )

  chat_msg = text_box.submit(query_message, [text_box,chatbot], [text_box, chatbot])
  bot_msg = chat_msg.then(chat_bot, chatbot, chatbot)
  bot_msg.then(lambda: gr.MultimodalTextbox(interactive=True), None, text_box)

if __name__ == "__main__":
    app.queue() 
    app.launch(share=True)