Update app_demo.py

app_demo.py

@@ -27,12 +27,6 @@ executor = ThreadPoolExecutor()
 model_cache = {}
 
 model_id = "Lykon/dreamshaper-xl-v2-turbo"
-#custom_pipe = DiffusionPipeline.from_pretrained(model_id, custom_pipeline="latent_consistency_txt2img", custom_revision="main")
-pipe = DiffusionPipeline.from_pretrained("SimianLuo/LCM_Dreamshaper_v7", custom_pipeline="latent_consistency_txt2img", custom_revision="main")
-pipe.to(torch_device="cpu", torch_dtype=DTYPE)
-pipe.safety_checker = None
-#pipe = StableDiffusionPipeline.from_pretrained(    model_id, safety_checker=None, torch_dtype=DTYPE, use_safetensors=True).to("cpu")
-
 custom_pipe = DiffusionPipeline.from_pretrained(
     model_id,
     custom_pipeline="latent_consistency_txt2img",
@@ -40,23 +34,22 @@ custom_pipe = DiffusionPipeline.from_pretrained(
     safety_checker=None,
     feature_extractor=None
 )
+custom_pipe.to(torch_device="cpu", torch_dtype=DTYPE)
+pipe = custom_pipe
 
 def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
     return random.randint(0, MAX_SEED) if randomize_seed else seed
 
-
 def save_image(img, profile: gr.OAuthProfile | None, metadata: dict):
     unique_name = str(uuid.uuid4()) + '.png'
     img.save(unique_name)
     gr_user_history.save_image(label=metadata["prompt"], image=img, profile=profile, metadata=metadata)
     return unique_name
 
-
 def save_images(image_array, profile: gr.OAuthProfile | None, metadata: dict):
     with ThreadPoolExecutor() as executor:
         return list(executor.map(save_image, image_array, [profile]*len(image_array), [metadata]*len(image_array)))
 
-
 def generate(prompt: str, seed: int = 0, width: int = 512, height: int = 512,
              guidance_scale: float = 8.0, num_inference_steps: int = 4,
              num_images: int = 1, randomize_seed: bool = False,
@@ -75,7 +68,6 @@ def generate(prompt: str, seed: int = 0, width: int = 512, height: int = 512,
                                                     "num_inference_steps": num_inference_steps})
     return paths, seed
 
-
 def validate_and_list_models(hfuser):
     try:
         models = api.list_models(author=hfuser)
@@ -83,7 +75,6 @@ def validate_and_list_models(hfuser):
     except Exception:
         return []
 
-
 def parse_user_model_dict(user_model_dict_str):
     try:
         data = ast.literal_eval(user_model_dict_str)
@@ -93,7 +84,6 @@ def parse_user_model_dict(user_model_dict_str):
     except Exception:
         return {}
 
-
 def load_model(model_id):
     if model_id in model_cache:
         return f"{model_id} loaded from cache"
@@ -104,14 +94,12 @@ def load_model(model_id):
     except Exception as e:
         return f"{model_id} failed to load: {str(e)}"
 
-
 def run_models(models, parallel):
     if parallel:
         futures = [executor.submit(load_model, m) for m in models]
         return [f.result() for f in futures]
     return [load_model(m) for m in models]
 
-
 with gr.Blocks() as demo:
     with gr.Row():
         gr.HTML("""
@@ -175,4 +163,4 @@ with gr.Blocks() as demo:
         fn=generate,
         inputs=[prompt, seed, width, height, guidance_scale, num_inference_steps, num_images, randomize_seed],
         outputs=[gallery, seed]
-    )
+    )