Update app.py

app.py

@@ -1,4 +1,4 @@
-import spaces
+import spaces  # Import spaces at the very beginning
 import os
 import pandas as pd
 import torch
@@ -15,8 +15,6 @@ from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5Tokeniz
 from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
 import gradio as gr
 from accelerate import Accelerator
-import huggingface_hub # Ensure this import is correct
-
 
 # Instantiate the Accelerator
 accelerator = Accelerator()
@@ -89,7 +87,7 @@ paused = False
 parsed_descriptions_queue = deque()
 
 @spaces.GPU
-def generate_and_store_descriptions(user_prompt, batch_size=100, max_iterations=50):
+def generate_descriptions(user_prompt, seed_words_input, batch_size=100, max_iterations=50):
     global paused
     descriptions = []
     description_queue = deque()
@@ -103,6 +101,9 @@ def generate_and_store_descriptions(user_prompt, batch_size=100, max_iterations=
     text_generator = pipeline('text-generation', model=model, tokenizer=tokenizer)
     print("Text generation pipeline initialized with 16-bit precision.")
 
+    # Populate the seed_words array with user input
+    seed_words.extend(re.findall(r'"(.*?)"', seed_words_input))
+
     while iteration_count < max_iterations:
         if paused:
             break
@@ -132,20 +133,18 @@ def generate_and_store_descriptions(user_prompt, batch_size=100, max_iterations=
             if iteration_count % 3 == 0:
                 parsed_descriptions = parse_descriptions(clean_description)
                 parsed_descriptions_queue.extend(parsed_descriptions)
-                # Return the parsed descriptions to update the Gradio UI
-                return list(parsed_descriptions_queue)
 
         iteration_count += 1
 
     return list(parsed_descriptions_queue)
 
 @spaces.GPU(duration=120)
-def generate_images_from_parsed_descriptions():
+def generate_images(parsed_descriptions):
     # If there are fewer than 13 descriptions, use whatever is available
-    if len(parsed_descriptions_queue) < 13:
-        prompts = list(parsed_descriptions_queue)
+    if len(parsed_descriptions) < 13:
+        prompts = parsed_descriptions
     else:
-        prompts = [parsed_descriptions_queue.popleft() for _ in range(13)]
+        prompts = [parsed_descriptions.pop(0) for _ in range(13)]
 
     # Generate images from the parsed descriptions
     images = []
@@ -155,16 +154,15 @@ def generate_images_from_parsed_descriptions():
     return images
 
 # Create Gradio Interface
-description_interface = gr.Interface(
-    fn=generate_and_store_descriptions,
-    inputs=gr.Textbox(lines=2, placeholder="Enter a prompt for descriptions..."),
-    outputs="json"
-)
+def combined_function(user_prompt, seed_words_input):
+    parsed_descriptions = generate_descriptions(user_prompt, seed_words_input)
+    images = generate_images(parsed_descriptions)
+    return images
 
-image_interface = gr.Interface(
-    fn=generate_images_from_parsed_descriptions,
-    inputs=None,
+interface = gr.Interface(
+    fn=combined_function,
+    inputs=[gr.Textbox(lines=2, placeholder="Enter a prompt for descriptions..."), gr.Textbox(lines=2, placeholder='Enter seed words in quotes, e.g., "cat", "dog", "sunset"...')],
     outputs=gr.Gallery()
 )
 
-gr.TabbedInterface([description_interface, image_interface], ["Generate Descriptions", "Generate Images"]).launch()
+interface.launch()