Update app.py

app.py

@@ -33,7 +33,7 @@ parsed_descriptions_queue = deque()
 
 # Usage limits
 MAX_DESCRIPTIONS = 30
-MAX_IMAGES = 3  # Limit to 3 images
+MAX_IMAGES = 4  # Limit to 4 images
 
 # Preload models and checkpoints
 print("Preloading models and checkpoints...")
@@ -90,7 +90,9 @@ def generate_description_prompt(user_prompt, text_generator):
     try:
         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
+        # Filter descriptions to ensure they are at least 4 words long
+        filtered_descriptions = [desc for desc in generated_descriptions if len(desc.split()) >= 4]
+        return filtered_descriptions if filtered_descriptions else None
     except Exception as e:
         print(f"Error generating descriptions: {e}")
         return None
@@ -100,22 +102,33 @@ def format_descriptions(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 = generate_description_prompt(user_prompt, text_generator)
-    if descriptions:
-        parsed_descriptions_queue.extend(descriptions)
+def generate_descriptions(user_prompt, seed_words_input, batch_size=100, max_iterations=2):  # Set max_iterations to 2
+    descriptions = []
+    for _ in range(2):  # Perform two iterations
+        new_descriptions = generate_description_prompt(user_prompt, text_generator)
+        if new_descriptions:
+            descriptions.extend(new_descriptions)
+            # Pick a random description to feed back into the seed bank for subject
+            random_description = random.choice(new_descriptions)
+            seed_words.append(random_description)
+
+    # Limit the number of descriptions to MAX_IMAGES (4)
+    if len(descriptions) > MAX_IMAGES:
+        descriptions = descriptions[:MAX_IMAGES]
+
+    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 MAX_IMAGES (3)
+def generate_images(parsed_descriptions, max_iterations=4):  # Set max_iterations to 4
+    # Limit the number of descriptions passed to the image generator to MAX_IMAGES (4)
     if len(parsed_descriptions) > MAX_IMAGES:
         parsed_descriptions = parsed_descriptions[:MAX_IMAGES]
 
     images = []
     for prompt in parsed_descriptions:
         try:
-            images.extend(pipe(prompt, num_inference_steps=4, height=1024, width=1024).images)  # Set resolution to 512 x 512
+            images.extend(pipe(prompt, num_inference_steps=4, height=1024, width=1024).images)  # Set resolution to 1024 x 1024
         except Exception as e:
             print(f"Error generating image for prompt '{prompt}': {e}")
 
@@ -140,4 +153,4 @@ if __name__ == '__main__':
         allow_flagging='never'  # Disable flagging
     )
 
-    interface.launch(share=True)
+    interface.launch(share=True)