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,