update

app.py

@@ -14,26 +14,44 @@ from scheduling_dpmsolver_multistep_inject import DPMSolverMultistepSchedulerInj
 from transformers import AutoProcessor, BlipForConditionalGeneration
 from share_btn import community_icon_html, loading_icon_html, share_js
 
-# load pipelines
-# sd_model_id = "runwayml/stable-diffusion-v1-5"
-sd_model_id = "stabilityai/stable-diffusion-2-1-base"
-#device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-torch_dtype = torch.float16
-if torch.cuda.is_available():
-  device = "cuda"
-elif torch.backends.mps.is_available():
-  device = "mps"
-  torch_dtype = torch.float32
-else:
-  device = "cpu"
-vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse", torch_dtype=torch_dtype)
-pipe = SemanticStableDiffusionImg2ImgPipeline_DPMSolver.from_pretrained(sd_model_id,vae=vae,torch_dtype=torch_dtype, safety_checker=None, requires_safety_checker=False).to(device)
-pipe.scheduler = DPMSolverMultistepSchedulerInject.from_pretrained(sd_model_id, subfolder="scheduler"
-                                                             , algorithm_type="sde-dpmsolver++", solver_order=2)
-
-blip_processor = AutoProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
-blip_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base",torch_dtype=torch_dtype).to(device)
+def load_models():
+    sd_model_id = "stabilityai/stable-diffusion-2-1-base"
+
+    torch_dtype = torch.float16
+    if torch.cuda.is_available():
+        device = "cuda"
+    elif torch.backends.mps.is_available():
+        device = "mps"
+        torch_dtype = torch.float32
+    else:
+        device = "cpu"
+    
+    vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse", torch_dtype=torch_dtype)
+    
+    pipe = SemanticStableDiffusionImg2ImgPipeline_DPMSolver.from_pretrained(
+        sd_model_id,
+        vae=vae,
+        torch_dtype=torch_dtype,
+        safety_checker=None,
+        requires_safety_checker=False
+    ).to(device)
+    
+    pipe.scheduler = DPMSolverMultistepSchedulerInject.from_pretrained(
+        sd_model_id,
+        subfolder="scheduler",
+        algorithm_type="sde-dpmsolver++",
+        solver_order=2
+    )
+    
+    blip_processor = AutoProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
+    blip_model = BlipForConditionalGeneration.from_pretrained(
+        "Salesforce/blip-image-captioning-base",
+        torch_dtype=torch_dtype
+    ).to(device)
+    
+    return vae, pipe, blip_processor, blip_model, device, torch_dtype
+
+vae, pipe, blip_processor, blip_model, device, torch_dtype = load_models()
 
 ## IMAGE CPATIONING ##
 def caption_image(input_image):
@@ -860,14 +878,19 @@ with gr.Blocks(css="style.css") as demo:
     clear_button.click(lambda: clear_components_output_vals, outputs = clear_components)
 
     reconstruct_button.click(lambda: ddpm_edited_image.update(visible=True), outputs=[ddpm_edited_image]).then(fn = reconstruct,
-                inputs = [tar_prompt,
-                image_caption,
-                tar_cfg_scale,
-                skip,
-                wts, zs,
-                do_reconstruction,
-                reconstruction,
-                          reconstruct_button],
+                inputs = [
+                    tar_prompt,
+                    image_caption,
+                    tar_cfg_scale,
+                    skip,
+                    wts,
+                    zs,
+                    attention_store,
+                    text_cross_attention_maps,
+                    do_reconstruction,
+                    reconstruction,
+                    reconstruct_button
+                ],
                 outputs = [ddpm_edited_image,reconstruction, ddpm_edited_image, do_reconstruction, reconstruct_button])
 
     randomize_seed.change(
@@ -905,8 +928,8 @@ with gr.Blocks(css="style.css") as demo:
                 threshold_3,
                     seed
                ],
-        outputs=[share_btn_container, box1, concept_1, guidnace_scale_1,neg_guidance_1, row1, row2,box2, concept_2, guidnace_scale_2,neg_guidance_2,row2, row3,sega_concepts_counter],
-        cache_examples=True
+        outputs=[share_btn_container, box1, concept_1, guidnace_scale_1,neg_guidance_1, row1, row2, box2, concept_2, guidnace_scale_2,neg_guidance_2,row2, row3,sega_concepts_counter],
+        cache_examples=False
     )
 
 demo.queue(default_concurrency_limit=1)