Spaces:
Running
Running
File size: 3,543 Bytes
64ceaa0 0d0b13c 64ceaa0 0d0b13c f6ef6ff 64ceaa0 0d0b13c 64ceaa0 1ceb41f c1ddcd2 64ceaa0 c1ddcd2 64ceaa0 c1ddcd2 64ceaa0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 |
import gradio as gr
import torch
import os
import gc
import random
from huggingface_hub import snapshot_download, HfApi
from diffusers import StableDiffusionXLPipeline, LCMScheduler
from PIL import Image
newcv = os.getenv("newcv")
newhf = os.getenv("newhf")
api = HfApi(token=newhf)
os.environ["XDG_CACHE_HOME"] = "/home/user/.cache"
os.environ["TRANSFORMERS_CACHE"] = "/home/user/.cache/huggingface/transformers"
os.environ["HF_HOME"] = "/home/user/.cache/huggingface"
models = [
"Niggendar/fastPhotoPony_v80MixB",
"Niggendar/realisticPonyPhoto_v10",
"Niggendar/realmix_v10",
"Niggendar/realmixpony_v01",
"Niggendar/realmixpony_v02",
"Niggendar/recondiff_v10",
"Niggendar/Regro",
"Niggendar/relhCheckpoint_v20",
]
loras = ["openskyml/lcm-lora-sdxl-turbo"]
pipe = None
cached = {}
cached_loras = {}
def get_lora(lora_id):
if lora_id in cached_loras:
return cached_loras[lora_id]
lora_dir = snapshot_download(repo_id=lora_id, use_auth_token=newhf, local_files_only=False ,allow_patterns=["*.safetensors", "*.bin"])
lora_files = [f for f in os.listdir(lora_dir) if f.endswith((".safetensors", ".bin"))]
lora_path = os.path.join(lora_dir, lora_files[0])
cached_loras[lora_id] = lora_path
return lora_path
def load_pipe(model_id, lora_id):
global pipe
if (model_id, lora_id) in cached:
pipe = cached[(model_id, lora_id)]
return
if pipe is not None:
pipe.to("meta")
pipe.unet = None
pipe.vae = None
pipe.text_encoder = None
del pipe
gc.collect()
cached.clear()
pipe = StableDiffusionXLPipeline.from_pretrained(model_id,torch_dtype=torch.float32,low_cpu_mem_usage=True )
pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
pipe.load_lora_weights(get_lora(lora_id))
pipe.to("cpu", dtype=torch.float32)
pipe.enable_attention_slicing()
cached[(model_id, lora_id)] = pipe
return gr.update(value='-')
def infer(model_id, lora_id, prompt, seed=None, steps=4, guid=0.1):
if seed is None or seed == "":
seed = random.randint(0, 2**32 - 1)
yield Image.new("RGB", (512, 512), color="gray"), gr.update(value='-')
image = pipe( prompt, generator=torch.manual_seed(int(seed)), num_inference_steps=steps,
guidance_scale=guid,width=128+256, height=128+256, added_cond_kwargs={} ).images[0]
yield image, gr.update(value='-')
with gr.Blocks() as demo:
with gr.Row():
with gr.Column(scale=2):
modeldrop=gr.Dropdown(models, label="Model",container=False)
loradrop=gr.Dropdown(loras, label="LCM LoRA",container=False,interactive=False,visible=False)
text= gr.Textbox(label="Prompt",container=False,placeholder="Prompt",value='')
gbtn=gr.Button(value="Generate")
with gr.Accordion(label="Settings", open=False):
seed=gr.Textbox(label="Seed",visible=False)
steps=gr.Slider(1, 15, value=4, step=1, label="Steps")
guidance=gr.Slider(0.0, 2.0, value=0.1, step=0.1, label="Guidance Scale")
with gr.Column(scale=3):
text2=gr.Textbox(label="Time",placeholder="timer",container=False,value='-')
imageout=gr.Image()
modeldrop.change(fn=load_pipe, inputs=[ modeldrop, loradrop ], outputs=[text2])
gbtn.click(fn=infer, inputs=[ modeldrop, loradrop, text, seed,steps, guidance ], outputs=[imageout,text2])
demo.queue()
demo.launch(server_name="0.0.0.0", server_port=7860)
|