Update app.py

app.py

@@ -21,14 +21,14 @@ MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
 #@spaces.GPU #[uncomment to use ZeroGPU]
-def infer_t2i(model_repo_id, prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, progress=gr.Progress(track_tqdm=True)):
+def infer_t2i(model, prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, progress=gr.Progress(track_tqdm=True)):
 
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
         
     generator = torch.Generator().manual_seed(seed)
 
-    pipe = StableDiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
+    pipe = StableDiffusionPipeline.from_pretrained(all_model_id[model], torch_dtype=torch_dtype)
     pipe = pipe.to(device)
     
     image = pipe(
@@ -44,14 +44,14 @@ def infer_t2i(model_repo_id, prompt, negative_prompt, seed, randomize_seed, widt
     return image, seed
 
 #@spaces.GPU #[uncomment to use ZeroGPU]
-def infer_i2i(model_repo_id, prompt, image, strength, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, progress=gr.Progress(track_tqdm=True)):
+def infer_i2i(model, prompt, image, strength, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, progress=gr.Progress(track_tqdm=True)):
 
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
         
     generator = torch.Generator().manual_seed(seed)
 
-    pipe = StableDiffusionImg2ImgPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
+    pipe = StableDiffusionImg2ImgPipeline.from_pretrained(all_model_id[model], torch_dtype=torch_dtype)
     pipe = pipe.to(device)
     
     image = pipe(
@@ -69,14 +69,14 @@ def infer_i2i(model_repo_id, prompt, image, strength, negative_prompt, seed, ran
     return image, seed
 
 #@spaces.GPU #[uncomment to use ZeroGPU]
-def infer_ip_adapter(model_repo_id, prompt, image, scale, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, progress=gr.Progress(track_tqdm=True)):
+def infer_ip_adapter(model, prompt, image, scale, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, progress=gr.Progress(track_tqdm=True)):
 
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
         
     generator = torch.Generator().manual_seed(seed)
 
-    pipe = StableDiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
+    pipe = StableDiffusionPipeline.from_pretrained(all_model_id[model], torch_dtype=torch_dtype)
     pipe = pipe.to(device)
     pipeline.load_ip_adapter("h94/IP-Adapter", subfolder="models", weight_name="ip-adapter-plus_sd15.bin")
     pipeline.set_ip_adapter_scale(scale)
@@ -196,7 +196,7 @@ with gr.Blocks(css=css) as demo:
     
             run_button.click(
                 fn=infer_t2i,
-                inputs = [all_model_id[str(model_choice)], prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
+                inputs = [prompt, model_choice, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
                 outputs = [result, seed]
             )
 
@@ -295,7 +295,7 @@ with gr.Blocks(css=css) as demo:
     
             run_button.click(
                 fn=infer_i2i,
-                inputs = [all_model_id[str(model_choice)], prompt, image_upload_input, editing_strength, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
+                inputs = [prompt, model_choice, image_upload_input, editing_strength, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
                 outputs = [result, seed]
             )
 
@@ -394,7 +394,7 @@ with gr.Blocks(css=css) as demo:
     
             run_button.click(
                 fn=infer_ip_adapter,
-                inputs = [all_model_id[str(model_choice)], prompt, image_upload_input, ip_adapter_scale, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
+                inputs = [prompt, model_choice, image_upload_input, ip_adapter_scale, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
                 outputs = [result, seed]
             )