Update app.py

app.py

@@ -10,7 +10,8 @@ from src.linfusion import LinFusion
 device = "cuda" if torch.cuda.is_available() else "cpu"
 all_model_id = {
     "DreamShaper-8": "Lykon/dreamshaper-8",
-    "RealisticVision-v4.0": "SG161222/Realistic_Vision_V4.0_noVAE"
+    "RealisticVision-v4.0": "SG161222/Realistic_Vision_V4.0_noVAE",
+    "SD-v1.4": "CompVis/stable-diffusion-v1-4"
 }
 
 if torch.cuda.is_available():
@@ -20,6 +21,18 @@ else:
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
+pipes = {}
+for model_id, repo_id in all_model_id.items():
+    pipes[model_id + '_t2i'] = StableDiffusionPipeline.from_pretrained(repo_id, torch_dtype=torch_dtype)
+    LinFusion.construct_for(pipes[model_id + '_t2i'])
+    
+    pipes[model_id + '_ip_adapter'] = StableDiffusionPipeline.from_pretrained(repo_id, torch_dtype=torch_dtype)
+    pipes[model_id + '_ip_adapter'].load_ip_adapter("h94/IP-Adapter", subfolder="models", weight_name="ip-adapter-plus_sd15.bin")
+    pipes[model_id + '_ip_adapter'].set_ip_adapter_scale(scale)
+    LinFusion.construct_for(pipes[model_id + '_ip_adapter'])
+    
+    pipes[model_id + '_i2i'] = StableDiffusionImg2ImgPipeline.from_pretrained(repo_id, torch_dtype=torch_dtype)
+    LinFusion.construct_for(pipes[model_id + '_i2i'])
 
 @spaces.GPU
 def infer_t2i(model, prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, progress=gr.Progress(track_tqdm=True)):
@@ -29,9 +42,7 @@ def infer_t2i(model, prompt, negative_prompt, seed, randomize_seed, width, heigh
         
     generator = torch.Generator().manual_seed(seed)
 
-    pipe = StableDiffusionPipeline.from_pretrained(all_model_id[model], torch_dtype=torch_dtype)
-    pipe = pipe.to(device)
-    linfusion = LinFusion.construct_for(pipe)
+    pipe = pipes[model + '_t2i'].to(device)
     
     image = pipe(
         prompt = prompt, 
@@ -53,9 +64,7 @@ def infer_i2i(model, prompt, image, strength, negative_prompt, seed, randomize_s
         
     generator = torch.Generator().manual_seed(seed)
 
-    pipe = StableDiffusionImg2ImgPipeline.from_pretrained(all_model_id[model], torch_dtype=torch_dtype)
-    pipe = pipe.to(device)
-    linfusion = LinFusion.construct_for(pipe)
+    pipe = pipes[model + '_i2i'].to(device)
     
     image = pipe(
         prompt = prompt, 
@@ -79,11 +88,7 @@ def infer_ip_adapter(model, prompt, image, scale, negative_prompt, seed, randomi
         
     generator = torch.Generator().manual_seed(seed)
 
-    pipe = StableDiffusionPipeline.from_pretrained(all_model_id[model], torch_dtype=torch_dtype)
-    pipe = pipe.to(device)
-    pipe.load_ip_adapter("h94/IP-Adapter", subfolder="models", weight_name="ip-adapter-plus_sd15.bin")
-    pipe.set_ip_adapter_scale(scale)
-    linfusion = LinFusion.construct_for(pipe)
+    pipe = pipes[model + 'ip_adapter'].to(device)
     
     image = pipe(
         prompt = prompt,