Update app.py

app.py

@@ -83,68 +83,41 @@ def initialize_diffusers():
 pipe = initialize_diffusers()
 print("Models and checkpoints preloaded.")
 
-def generate_description_prompt(subject, user_prompt, text_generator):
-    prompt = f"write concise vivid visual description enclosed in brackets like [ <description> ] less than 100 words of {user_prompt} different from {subject}. "
+def generate_description_prompt(user_prompt, text_generator):
+    injected_prompt = f"write three concise descriptions enclosed in brackets like [ <description> ] less than 100 words each of {user_prompt}. "
+    max_length = 110  # Set max token length to 110
+
     try:
-        generated_text = text_generator(prompt, max_length=160, num_return_sequences=1, truncation=True)[0]['generated_text']
-        generated_description = re.sub(rf'{re.escape(prompt)}\s*', '', generated_text).strip()  # Remove the prompt from the generated text
-        return generated_description if generated_description else None
+        generated_text = text_generator(injected_prompt, max_length=max_length, num_return_sequences=1, truncation=True)[0]['generated_text']
+        generated_descriptions = re.findall(r'\[([^\[\]]+)\]', generated_text)  # Extract descriptions enclosed in brackets
+        return generated_descriptions if generated_descriptions else None
     except Exception as e:
-        print(f"Error generating description for subject '{subject}': {e}")
+        print(f"Error generating descriptions: {e}")
         return None
 
-def parse_descriptions(text):
-    descriptions = re.findall(r'\[([^\[\]]+)\]', text)
-    descriptions = [desc.strip() for desc in descriptions if len(desc.split()) >= 3]
-    return descriptions
-
 def format_descriptions(descriptions):
     formatted_descriptions = "\n".join(descriptions)
     return formatted_descriptions
 
 @spaces.GPU
 def generate_descriptions(user_prompt, seed_words_input, batch_size=100, max_iterations=1):  # Set max_iterations to 1
-    descriptions = []
-    description_queue = deque()
-    iteration_count = 0
-
-    seed_words.extend(re.findall(r'"(.*?)"', seed_words_input))
-
-    for _ in range(2):  # Perform two iterations
-        while iteration_count < max_iterations and len(parsed_descriptions_queue) < MAX_DESCRIPTIONS:
-            available_subjects = [word for word in seed_words if word not in used_words]
-            if not available_subjects:
-                print("No more available subjects to use.")
-                break
-
-            subject = random.choice(available_subjects)
-            generated_description = generate_description_prompt(subject, user_prompt, text_generator)
-
-            if generated_description:
-                clean_description = generated_description.encode('ascii', 'ignore').decode('ascii')
-                description_queue.append({'subject': subject, 'description': clean_description})
-
-                print(f"Generated description for subject '{subject}': {clean_description}")
-
-                used_words.add(subject)
-                seed_words.append(clean_description)
-
-                parsed_descriptions = parse_descriptions(clean_description)
-                parsed_descriptions_queue.extend(parsed_descriptions)
-
-            iteration_count += 1
-
-    return list(parsed_descriptions_queue)
+    descriptions = generate_description_prompt(user_prompt, text_generator)
+    if descriptions:
+        parsed_descriptions_queue.extend(descriptions)
+    return list(parsed_descriptions_queue)[:MAX_IMAGES]
 
 @spaces.GPU(duration=120)
 def generate_images(parsed_descriptions, max_iterations=3):  # Set max_iterations to 3
-    # Limit the number of descriptions passed to the image generator to 2
+    # Limit the number of descriptions passed to the image generator to MAX_IMAGES (3)
     if len(parsed_descriptions) > MAX_IMAGES:
         parsed_descriptions = parsed_descriptions[:MAX_IMAGES]
 
     images = []
     for prompt in parsed_descriptions:
-        images.extend(pipe(prompt, num_inference_steps=max_iterations, height=512, width=512).images)  # Set resolution to 512 x 512
+        try:
+            images.extend(pipe(prompt, num_inference_steps=4, height=1024, width=1024).images)  # Set resolution to 512 x 512
+        except Exception as e:
+            print(f"Error generating image for prompt '{prompt}': {e}")
 
     return images
 
@@ -161,10 +134,10 @@ if __name__ == '__main__':
 
     interface = gr.Interface(
         fn=generate_and_display,
-        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"...')],
+        inputs=[gr.Textbox(lines=2, placeholder="Enter a prompt for descriptions..."), gr.Textbox(lines=2, placeholder='Enter example in quotes, e.g., "cat", "dog", "sunset"...')],
         outputs=[gr.Textbox(label="Generated Descriptions"), gr.Gallery(label="Generated Images")],
         live=False,  # Set live to False
         allow_flagging='never'  # Disable flagging
     )
 
-    interface.launch(share=True)
+    interface.launch(share=True)