Update app.py

app.py

@@ -1,126 +1,109 @@
 import gradio as gr
 from random import randint
-from all_models import models
 from datetime import datetime
 
-kii = "mohawk femboy racecar driver"
-
-
-model_group_selector = gr.Dropdown(
-    choices=list(model_groups.keys()),
-    label="Select Model Group",
-    value="Group A",  # Default group
-)
-
-# Update output based on the selected group
-def update_models(group_name):
-    return model_groups[group_name]  # Return the list of models for the selected group
-
-# Connect the dropdown to the update function
-model_group_selector.change(update_models, inputs=[model_group_selector], outputs=[...])
-
-
-def get_current_time():
-    now = datetime.now()
-    current_time = now.strftime("%Y-%m-%d %H:%M:%S")
-    return f'{kii} {current_time}'
-
-
+# Define manual grouping for models
+model_groups = {
+    "Group A": [
+        "Bakanayatsu/ponyDiffusion-V6-XL-Turbo-DPO",
+        "John6666/photo-realistic-pony-v5-sdxl",
+        "John6666/photo-realistic-pony-v5-sdxl",
+    ],
+    "Group B": [
+        "John6666/jib-mix-pony-realistic-v2-sdxl",
+        "John6666/3x3x3mixxl-v2-sdxl-spo",
+        "John6666/3x3x3mixxl-v2-sdxl",
+    ],
+    "Group C": [
+        "John6666/3x3x3mixxl-v2-sdxl-spo",
+        "John6666/3x3x3mixxl-v2-sdxl",
+        "John6666/titania-mix-realistic-pony-gbv10-sdxl",
+    ],
+    "Group D": [
+        "John6666/titania-mix-realistic-pony-gbv20-sdxl",
+        "John6666/titania-mix-realistic-pony-gbv30-sdxl",
+        "John6666/mala-anime-mix-nsfw-pony-xl-v3-sdxl",
+    ],
+}
+
+# Placeholder for models
+models_load = {model: None for group in model_groups.values() for model in group}
+
+# Function to simulate model loading
 def load_fn(models):
-    models_load = {}
+    loaded_models = {}
     for model in models:
-        if model not in models_load:
-            try:
-                m = gr.load(f'models/{model}')  # Adjust `gr.load` as needed
-            except Exception as error:
-                m = gr.Interface(lambda txt: None, ['text'], ['image'])
-            models_load[model] = m
-    return models_load
-
-
-models_load = load_fn(models)
-
-def extend_choices(choices, num_models):
-    return choices + (num_models - len(choices)) * ['NA']
-
-
-def update_imgbox(choices, num_models):
-    choices_plus = extend_choices(choices, num_models)
-    return [gr.Image(None, label=m, visible=(m != 'NA')) for m in choices_plus]
+        try:
+            # Simulate model loading
+            loaded_models[model] = lambda txt: f"Generated image for {txt}"
+        except Exception as error:
+            loaded_models[model] = lambda txt: None
+    return loaded_models
 
+models_load = load_fn(models_load.keys())
 
+# Function to simulate generation with tally tracking
 def gen_fn(model_str, prompt, tallies):
     if model_str == 'NA':
         return None, tallies
 
     noise = str(randint(0, 9999))
     combined_prompt = f'{prompt} {model_str} {noise}'
-    print(f"Generating with prompt: {combined_prompt}")  # Debug line
+    print(f"Generating with prompt: {combined_prompt}")
 
     try:
-        result = models_load.get(model_str, lambda txt: None)(combined_prompt)
-        if result is not None:
-            tallies[model_str] += 1
+        result = models_load[model_str](combined_prompt)
+        tallies[model_str] += 1
         return result, tallies
     except Exception as e:
         print(f"Error generating for {model_str}: {e}")
         return None, tallies
 
-
+# Build the Gradio interface
 def make_me():
     with gr.Row():
-        # Input elements
         model_group_selector = gr.Dropdown(
-            choices=list(model_groups.keys()), label="Select Model Group", value="Group A"
+            choices=list(model_groups.keys()),
+            label="Select Model Group",
+            value="Group A",  # Default group
         )
-        txt_input = gr.Textbox(label='Your prompt:', lines=2, value=kii)
+        txt_input = gr.Textbox(label='Your prompt:', lines=2, value="mohawk femboy racecar driver")
         gen_button = gr.Button('Generate images', elem_id="generate-btn")
-        stop_button = gr.Button('Stop', variant='secondary', interactive=False)
-
-        gr.HTML("""
-        <div style="text-align: center; max-width: 100%; margin: 0 auto;">
-            <body></body>
-        </div>
-        """)
 
     with gr.Row():
-        # Output elements
-        output = gr.State([])
-        tally_boxes = gr.State({})
-        output_display = gr.Column()
-        with output_display:
-            result_images = []
-            tally_counters = []
+        output = gr.State([])  # Placeholder for output elements
+        tallies = gr.State({})  # Tally counter for each model
+        result_images = gr.Column()
+        tally_counters = gr.Column()
 
         def update_outputs(group_name):
             selected_models = model_groups[group_name]
-            result_images.clear()
-            tally_counters.clear()
-            for model in selected_models:
-                result_images.append(gr.Image(label=model, width=170, height=170))
-                tally_counters.append(gr.Textbox(value="0", label=f"Tally for {model}", interactive=False))
-            return result_images, tally_counters, {model: 0 for model in selected_models}
+            outputs = [gr.Image(label=model, visible=True, width=170, height=170) for model in selected_models]
+            tallies_dict = {model: 0 for model in selected_models}
+            tally_boxes = [gr.Textbox(value="0", label=f"Tally for {model}", interactive=False) for model in selected_models]
+            return outputs, tally_boxes, tallies_dict
 
         model_group_selector.change(
-            update_outputs, [model_group_selector], [output, tally_boxes]
+            update_outputs,
+            inputs=[model_group_selector],
+            outputs=[result_images, tally_counters, tallies],
         )
 
         def generate_images(prompt, outputs, tallies):
-            for idx, model_element in enumerate(outputs):
+            for idx, model_output in enumerate(outputs):
                 model_str = list(tallies.keys())[idx]
                 result, tallies = gen_fn(model_str, prompt, tallies)
-                model_element.update(value=result)
-            for idx, tally_box in enumerate(tally_counters):
-                tally_box.update(value=str(tallies[list(tallies.keys())[idx]]))
+                model_output.update(value=result)
+            for idx, tally_box in enumerate(tallies.values()):
+                tally_counters[idx].update(value=str(tally_box))
             return tallies
 
         gen_button.click(
             generate_images,
-            inputs=[txt_input, output, tally_boxes],
-            outputs=[tally_boxes],
+            inputs=[txt_input, output, tallies],
+            outputs=[tallies],
         )
 
-
 js_code = """
 <script>
     const originalScroll = window.scrollTo;
@@ -136,7 +119,6 @@ js_code = """
 """
 
 with gr.Blocks() as demo:
-    gr.Markdown("<div></div>")
     make_me()
     gr.Markdown(js_code)