Update app.py

app.py

@@ -1,149 +1,96 @@
 import gradio as gr
 import numpy as np
 import random
+import matplotlib.pyplot as plt
+import matplotlib.image as mpimg
+from embedding import model, docs_list, doc_embeddings, embedded_dict
+import openai
 from openai import OpenAI
-
 import os
 
 api_key = os.getenv("YOUR_OPENAI_KEY")
 client = OpenAI(api_key=api_key)
 
-#import spaces #[uncomment to use ZeroGPU]
-from diffusers import DiffusionPipeline
-import torch
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model_repo_id = "stabilityai/sdxl-turbo" #Replace to the model you would like to use
-
-if torch.cuda.is_available():
-    torch_dtype = torch.float16
-else:
-    torch_dtype = torch.float32
-
-pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
-pipe = pipe.to(device)
-
-MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 1024
+# Function to generate meme caption using GPT model
+def generate_meme_caption(topic):
+    response = openai.ChatCompletion.create(
+        model="gpt-4o-mini",
+        messages=[
+            {
+                "role": "user",
+                "content": f'''You are a loser, a meme lover who is always on the internet. 
+                You are a meme professional, so you will generate short captions for memes based on the following rules:
+                Do not add emoji.
+                Do not add any punctuation.
+                Must be funny.
+                Only if some topic related to friends, use "bro" or "blud".
+                Answer from previous question can not be the same for this time prompt.
+                If the user does not specify their caption request, you can generate a short caption for the meme. 
+                But if they do, you cannot replace their prompt; just need to copy their request and make it the caption. 
+                Now generate a caption based on this topic: {topic}.
+                '''
+            }
+        ],
+        max_tokens=250,
+        temperature=1.1
+    )
 
-#@spaces.GPU #[uncomment to use ZeroGPU]
-def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, progress=gr.Progress(track_tqdm=True)):
+    caption = response['choices'][0]['message']['content']
+    return caption
 
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-        
-    generator = torch.Generator().manual_seed(seed)
+# Function to create meme with caption
+def create_meme(prompt):
+    # Tạo embedding cho prompt
+    prompt_embedding = model.encode([prompt])
     
-    image = pipe(
-        prompt = prompt, 
-        negative_prompt = negative_prompt,
-        guidance_scale = guidance_scale, 
-        num_inference_steps = num_inference_steps, 
-        width = width, 
-        height = height,
-        generator = generator
-    ).images[0] 
+    # Tính toán độ tương đồng giữa prompt và các văn bản
+    similarities = model.similarity(prompt_embedding, doc_embeddings)
     
-    return image, seed
-
-examples = [
-    "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
-    "An astronaut riding a green horse",
-    "A delicious ceviche cheesecake slice",
-]
-
-css="""
-#col-container {
-    margin: 0 auto;
-    max-width: 640px;
-}
-"""
-
-with gr.Blocks(css=css) as demo:
+    # Lấy chỉ số văn bản có độ tương đồng cao nhất
+    closest_text_idx = np.argmax(similarities)
     
-    with gr.Column(elem_id="col-container"):
-        gr.Markdown(f"""
-        # Text-to-Image Gradio Template
-        """)
+    # Tải ảnh meme tương ứng
+    img_path = f'/kaggle/input/memedata/{embedded_dict[closest_text_idx]["filename"]}'
+    img = mpimg.imread(img_path)
+    
+    # Tạo khoảng trắng để hiển thị caption
+    white_space_height = 100
+    white_space = np.ones((white_space_height, img.shape[1], img.shape[2]), dtype=img.dtype) * 255
+    
+    # Kết hợp khoảng trắng với hình ảnh gốc
+    combined_img = np.vstack((white_space, img))
+    
+    # Sinh caption từ GPT model
+    caption = generate_meme_caption(prompt)
+    
+    # Hiển thị ảnh với caption
+    plt.imshow(combined_img)
+    plt.axis('off')
+    plt.text(
+        x=combined_img.shape[1] / 2, 
+        y=white_space_height / 2, 
+        s=caption, 
+        fontsize=11, color='black', ha='center', va='top', fontweight='bold'
+    )
+    
+    plt.show()
+
+# Gradio interface
+with gr.Blocks() as demo:
+    with gr.Column():
+        gr.Markdown("# Meme Generator using Embeddings and GPT")
         
-        with gr.Row():
-            
-            prompt = gr.Text(
-                label="Prompt",
-                show_label=False,
-                max_lines=1,
-                placeholder="Enter your prompt",
-                container=False,
-            )
-            
-            run_button = gr.Button("Run", scale=0)
+        prompt = gr.Textbox(label="Enter your meme prompt:")
+        run_button = gr.Button("Generate Meme")
         
-        result = gr.Image(label="Result", show_label=False)
-
-        with gr.Accordion("Advanced Settings", open=False):
-            
-            negative_prompt = gr.Text(
-                label="Negative prompt",
-                max_lines=1,
-                placeholder="Enter a negative prompt",
-                visible=False,
-            )
-            
-            seed = gr.Slider(
-                label="Seed",
-                minimum=0,
-                maximum=MAX_SEED,
-                step=1,
-                value=0,
-            )
-            
-            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-            
-            with gr.Row():
-                
-                width = gr.Slider(
-                    label="Width",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024, #Replace with defaults that work for your model
-                )
-                
-                height = gr.Slider(
-                    label="Height",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024, #Replace with defaults that work for your model
-                )
-            
-            with gr.Row():
-                
-                guidance_scale = gr.Slider(
-                    label="Guidance scale",
-                    minimum=0.0,
-                    maximum=10.0,
-                    step=0.1,
-                    value=0.0, #Replace with defaults that work for your model
-                )
-                
-                num_inference_steps = gr.Slider(
-                    label="Number of inference steps",
-                    minimum=1,
-                    maximum=50,
-                    step=1,
-                    value=2, #Replace with defaults that work for your model
-                )
+        result = gr.Plot(label="Generated Meme")
         
-        gr.Examples(
-            examples = examples,
-            inputs = [prompt]
-        )
-    gr.on(
-        triggers=[run_button.click, prompt.submit],
-        fn = infer,
-        inputs = [prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
-        outputs = [result, seed]
-    )
+        def gradio_infer(prompt):
+            plt.figure(figsize=(10, 8))
+            create_meme(prompt)
+            plt.savefig('/tmp/meme.png')  # Lưu hình tạm thời
+            return plt
+        
+        run_button.click(gradio_infer, inputs=[prompt], outputs=[result])
 
-demo.queue().launch()
+demo.launch()