Update inference.py

inference.py

@@ -1,62 +1,3 @@
-from lcm_pipeline import LatentConsistencyModelPipeline
-from lcm_scheduler import LCMScheduler
+import gradio as gr
 
-from diffusers import AutoencoderKL, UNet2DConditionModel
-from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
-from transformers import CLIPTokenizer, CLIPTextModel, CLIPImageProcessor
-
-import os
-import torch
-from tqdm import tqdm
-from safetensors.torch import load_file
-
-# Input Prompt:
-prompt = "Self-portrait oil painting, a beautiful cyborg with golden hair"
-
-# Save Path:
-save_path = "./lcm_images"
-os.makedirs(save_path, exist_ok=True)
-
-
-# Origin SD Model ID:
-model_id = "digiplay/DreamShaper_7"
-
-
-# Initalize Diffusers Model:
-vae = AutoencoderKL.from_pretrained(model_id, subfolder="vae")
-text_encoder = CLIPTextModel.from_pretrained(model_id, subfolder="text_encoder")
-tokenizer = CLIPTokenizer.from_pretrained(model_id, subfolder="tokenizer")
-unet = UNet2DConditionModel.from_pretrained(model_id, subfolder="unet", device_map=None, low_cpu_mem_usage=False, local_files_only=True)
-safety_checker = StableDiffusionSafetyChecker.from_pretrained(model_id, subfolder="safety_checker")
-feature_extractor = CLIPImageProcessor.from_pretrained(model_id, subfolder="feature_extractor")
-
-
-# Initalize Scheduler:
-scheduler = LCMScheduler(beta_start=0.00085, beta_end=0.0120, beta_schedule="scaled_linear", prediction_type="epsilon")
-
-
-# Replace the unet with LCM:
-lcm_unet_ckpt = "./LCM_Dreamshaper_v7_4k.safetensors"
-ckpt = load_file(lcm_unet_ckpt)
-m, u = unet.load_state_dict(ckpt, strict=False)
-if len(m) > 0:
-    print("missing keys:")
-    print(m)
-if len(u) > 0:
-    print("unexpected keys:")
-    print(u)
-
-
-# LCM Pipeline:
-pipe = LatentConsistencyModelPipeline(vae=vae, text_encoder=text_encoder, tokenizer=tokenizer, unet=unet, scheduler=scheduler, safety_checker=safety_checker, feature_extractor=feature_extractor)
-pipe = pipe.to("cuda")
-
-
-# Output Images:
-images = pipe(prompt=prompt, num_images_per_prompt=4, num_inference_steps=4, guidance_scale=8.0, lcm_origin_steps=50).images
-
-# Save Images:
-for i in tqdm(range(len(images))):
-    output_path = os.path.join(save_path, "{}.png".format(i))
-    image = images[i]
-    image.save(output_path) 
+gr.load("models/SimianLuo/LCM_Dreamshaper_v7").launch()