Update app.py

app.py

@@ -1,54 +1,27 @@
 import gradio as gr
-import random
-from PIL import Image, ImageDraw
+import torch
+from transformers import GPT2Tokenizer, GPT2LMHeadModel
+import matplotlib.pyplot as plt
+import numpy as np
 
-def smooth_color(color, threshold):
-    if color == " ":
-        return random.choice(["FF0000", "32FF00", "00FFFF", "FFFFFF", "000000"])
-    return color
+# Load pre-trained GPT-2 model and tokenizer
+model = GPT2LMHeadModel.from_pretrained("gpt2")
+tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
 
+# Define a function to generate color based on text prompt
 def generate_color(prompt):
-    color_red = " "
-    color_green = " "
-    color_blue = " "
-    color_white = " "
-    color_none = " "
+    input_ids = tokenizer.encode(prompt, return_tensors='pt')
+    output = model.generate(input_ids, max_length=50, num_return_sequences=1, no_repeat_ngram_size=2)
+    color_name = tokenizer.decode(output[0], skip_special_tokens=True)
     
-    red = ['e', 'f', 'g', 'i', 'j', 'k', 'm', 'n', 'p', 'r']
-    green = ['c', 'd', 'h', 'l', 'o', 'q', 's', 'u', 'y', 'w']
-    blue = ['a', 'b', 't', 'v', 'x', 'z']
-    sym = ["!", "@", "#", "$", "%", "^", "&", "*", "(", ")", "<", ">", "/", "'", "-", "+", "="]
-
-    redded = []
-    greened = []
-    blued = []
-    whited = []
-    nothing = []
+    # Create an image with the generated color
+    color = [int(ord(char) * 255 / 122) for char in color_name[:3]]
+    img = np.full((100, 100, 3), color, dtype=np.uint8)
     
-    for letter in prompt:
-        if letter in red:
-            redded.append(letter)
-        elif letter in green:
-            greened.append(letter)
-        elif letter in blue:
-            blued.append(letter)
-        elif letter == " ":
-            whited.append(random.choice(sym))
-        else:
-            nothing.append(letter)
+    return img
 
-    color_red = smooth_color(color_red, len(redded))
-    color_green = smooth_color(color_green, len(greened))
-    color_blue = smooth_color(color_blue, len(blued))
-    color_white = smooth_color(color_white, len(whited))
-    color_none = smooth_color(color_none, len(nothing))
-    
-    # Convert the color to an image using PIL
-    image = Image.new('RGB', (100, 100), color="#"+color_red+color_green+color_blue)
-    draw = ImageDraw.Draw(image)
-    draw.text((20, 40), "".join(whited), fill="#"+color_white)
-    
-    return image
+# Create Gradio interface
+inputs = gr.Textbox(lines=2, label="Enter a text prompt (e.g., 'a color that represents happiness'):")
+output = gr.Image(type="numpy", label="Generated color:")
 
-iface = gr.Interface(generate_color, "text", "image", theme="soft")
-iface.launch()
+gr.Interface(fn=generate_color, inputs=inputs, outputs=output, title="AI Color Generator", description="Generate a color based on a text prompt using GPT-2 model.").launch()