Update app.py

app.py

@@ -32,7 +32,7 @@ parsed_descriptions_queue = deque()
 
 # Usage limits
 MAX_DESCRIPTIONS = 30
-MAX_IMAGES = 12
+MAX_IMAGES = 10
 
 def initialize_diffusers():
     from optimum.quanto import freeze, qfloat8, quantize
@@ -88,7 +88,7 @@ def parse_descriptions(text):
     return descriptions
 
 @spaces.GPU
-def generate_descriptions(user_prompt, seed_words_input, batch_size=100, max_iterations=50):
+def generate_descriptions(user_prompt, seed_words_input, batch_size=100, max_iterations=3):
     descriptions = []
     description_queue = deque()
     iteration_count = 0
@@ -102,36 +102,36 @@ def generate_descriptions(user_prompt, seed_words_input, batch_size=100, max_ite
 
     seed_words.extend(re.findall(r'"(.*?)"', seed_words_input))
 
-    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
+    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})
+            subject = random.choice(available_subjects)
+            generated_description = generate_description_prompt(subject, user_prompt, text_generator)
 
-            print(f"Generated description for subject '{subject}': {clean_description}")
+            if generated_description:
+                clean_description = generated_description.encode('ascii', 'ignore').decode('ascii')
+                description_queue.append({'subject': subject, 'description': clean_description})
 
-            used_words.add(subject)
-            seed_words.append(clean_description)
+                print(f"Generated description for subject '{subject}': {clean_description}")
+
+                used_words.add(subject)
+                seed_words.append(clean_description)
 
-            if iteration_count % 3 == 0:
                 parsed_descriptions = parse_descriptions(clean_description)
                 parsed_descriptions_queue.extend(parsed_descriptions)
 
-        iteration_count += 1
+            iteration_count += 1
 
     return list(parsed_descriptions_queue)
 
 @spaces.GPU(duration=120)
-def generate_images(parsed_descriptions):
+def generate_images(parsed_descriptions, max_iterations=10):
     pipe = initialize_diffusers()
-    
+
     if len(parsed_descriptions) < MAX_IMAGES:
         prompts = parsed_descriptions
     else:
@@ -139,20 +139,20 @@ def generate_images(parsed_descriptions):
 
     images = []
     for prompt in prompts:
-        images.extend(pipe(prompt, num_images=1).images)
+        images.extend(pipe(prompt, num_images=1, num_inference_steps=max_iterations, height=1024, width=1024).images)  # Define the resolution here
 
     return images
 
 def combined_function(user_prompt, seed_words_input):
     parsed_descriptions = generate_descriptions(user_prompt, seed_words_input)
     images = generate_images(parsed_descriptions)
-    return images
+    return parsed_descriptions, images
 
 if __name__ == '__main__':
     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()
+        outputs=[gr.Textbox(label="Generated Descriptions"), gr.Gallery(label="Generated Images")]
     )
 
     interface.launch()