Update app.py

app.py

@@ -17,8 +17,10 @@ def load_fn(models):
         if model not in models_load.keys():
             try:
                 m = gr.load(f'models/{model}')
+                print(f"Loaded model {model} with interface: {m}")
             except Exception as error:
                 m = gr.Interface(lambda txt: None, ['text'], ['image'])
+                print(f"Failed to load model {model}: {error}")
             models_load.update({model: m})
 
 load_fn(models)
@@ -33,18 +35,32 @@ def update_imgbox(choices):
     choices_plus = extend_choices(choices)
     return [gr.Image(None, label=m, visible=(m != 'NA'), elem_id="custom_image") for m in choices_plus]
 
-def gen_fn(model_str, prompt, negative_prompt, max_retries=10):
+def gen_fn(model_str, prompt, negative_prompt, guidance_scale, seed, clip_skip, scheduler, randomize_seed):
     if model_str == 'NA':
         return None
     
+    if randomize_seed:
+        seed = randint(0, 99999999)
+    
     retries = 0
-    while retries < max_retries:
+    while retries < 10:
         try:
-            noise = str(randint(0, 9999999))
-            if hasattr(models_load[model_str], 'negative_prompt'):
-                result = models_load[model_str](f'{prompt} {noise}', negative_prompt=negative_prompt)
-            else:
-                result = models_load[model_str](f'{prompt} {noise}')
+            noise = str(seed)
+            full_prompt = f"{prompt} {noise}"
+            
+            if negative_prompt:
+                full_prompt += f" --negative_prompt {negative_prompt}"
+            
+            if guidance_scale:
+                full_prompt += f" --guidance_scale {guidance_scale}"
+            
+            if clip_skip:
+                full_prompt += f" --clip_skip {clip_skip}"
+            
+            if scheduler:
+                full_prompt += f" --scheduler {scheduler}"
+            
+            result = models_load[model_str](full_prompt)
             return result
         except Exception as e:
             # Check for specific error messages or status codes
@@ -54,8 +70,8 @@ def gen_fn(model_str, prompt, negative_prompt, max_retries=10):
                 print(f"Error generating image: {e}")
             
             retries += 1
-            if retries >= max_retries:
-                raise Exception(f"Failed to generate image after {max_retries} retries.")
+            if retries >= 10:
+                raise Exception(f"Failed to generate image after 10 retries.")
     
     return None
 
@@ -77,11 +93,18 @@ def make_me():
                 gen_button.click(on_generate_click, inputs=None, outputs=[gen_button, stop_button])
                 stop_button.click(on_stop_click, inputs=None, outputs=[gen_button, stop_button])
     
+    with gr.Row():
+        guidance_scale = gr.Slider(minimum=0, maximum=20, step=0.1, label='Guidance Scale', value=7.5)
+        seed = gr.Slider(minimum=0, maximum=99999999, step=1, label='Seed', value=randint(0, 99999999))
+        clip_skip = gr.Slider(minimum=0, maximum=5, step=1, label='CLIP Skip', value=1)
+        scheduler = gr.Dropdown(['DDIM', 'PNDM', 'LMSDiscrete', 'EulerAncestralDiscrete', 'DPMSolverMultistep', 'HeunDiscrete', 'EulerDiscrete', 'DPMSolverSinglestep', 'DEISMultistep', 'UniPCMultistep'], label='Scheduler', value='DDIM')
+        randomize_seed = gr.Checkbox(label='Randomize Seed', value=False)
+    
     with gr.Row():
         output = [gr.Image(label=m, min_width=250, height=250, elem_id="custom_image") for m in default_models]
         current_models = [gr.Textbox(m, visible=False) for m in default_models]
         for m, o in zip(current_models, output):
-            gen_event = gen_button.click(gen_fn, [m, txt_input, negative_txt_input], o)
+            gen_event = gen_button.click(gen_fn, [m, txt_input, negative_txt_input, guidance_scale, seed, clip_skip, scheduler, randomize_seed], o)
             stop_button.click(on_stop_click, inputs=None, outputs=[gen_button, stop_button], cancels=[gen_event])
     
     with gr.Accordion('Model selection', elem_id="custom_accordion"):