Spaces:
Running
on
Zero
Running
on
Zero
import gradio as gr | |
import torch | |
import os | |
import glob | |
from datetime import datetime | |
from PIL import Image | |
from diffusers.utils import load_image | |
from diffusers import EulerDiscreteScheduler | |
from pipline_StableDiffusion_ConsistentID import ConsistentIDStableDiffusionPipeline | |
import spaces | |
zero = torch.Tensor([0]).cuda() | |
print(zero.device) # <-- 'cpu' π€ | |
# Gets the absolute path of the current script | |
script_directory = os.path.dirname(os.path.realpath(__file__)) | |
# download ConsistentID checkpoint to cache | |
consistentID_path = hf_hub_download(repo_id="JackAILab/ConsistentID", filename="ConsistentID-v1.bin", repo_type="model") | |
def process(inputImage,prompt,negative_prompt): | |
device = zero.device # "cuda" | |
base_model_path = "SG161222/Realistic_Vision_V6.0_B1_noVAE" | |
### Load base model | |
pipe = ConsistentIDStableDiffusionPipeline.from_pretrained( | |
base_model_path, | |
torch_dtype=torch.float16, | |
use_safetensors=True, | |
variant="fp16" | |
).to(device) | |
### Load consistentID_model checkpoint | |
pipe.load_ConsistentID_model( | |
os.path.dirname(consistentID_path), | |
subfolder="", | |
weight_name=os.path.basename(consistentID_path), | |
trigger_word="img", | |
) | |
pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config) | |
# hyper-parameter | |
select_images = load_image(Image.fromarray(inputImage)) | |
num_steps = 50 | |
merge_steps = 30 | |
if prompt == "": | |
prompt = "A man, in a forest, adventuring" | |
if negative_prompt == "": | |
negative_prompt = "monochrome, lowres, bad anatomy, worst quality, low quality, blurry" | |
#Extend Prompt | |
prompt = "cinematic photo," + prompt + ", 50mm photograph, half-length portrait, film, bokeh, professional, 4k, highly detailed" | |
negtive_prompt_group="((((ugly)))), (((duplicate))), ((morbid)), ((mutilated)), [out of frame], extra fingers, mutated hands, ((poorly drawn hands)), ((poorly drawn face)), (((mutation))), (((deformed))), ((ugly)), blurry, ((bad anatomy)), (((bad proportions))), ((extra limbs)), cloned face, (((disfigured))). out of frame, ugly, extra limbs, (bad anatomy), gross proportions, (malformed limbs), ((missing arms)), ((missing legs)), (((extra arms))), (((extra legs))), mutated hands, (fused fingers), (too many fingers), (((long neck)))" | |
negative_prompt = negative_prompt + negtive_prompt_group | |
seed = torch.randint(0, 1000, (1,)).item() | |
generator = torch.Generator(device=device).manual_seed(seed) | |
images = pipe( | |
prompt=prompt, | |
width=512, | |
height=512, | |
input_id_images=select_images, | |
negative_prompt=negative_prompt, | |
num_images_per_prompt=1, | |
num_inference_steps=num_steps, | |
start_merge_step=merge_steps, | |
generator=generator, | |
).images[0] | |
current_date = datetime.today() | |
output_dir = script_directory + f"/images/gradio_outputs" | |
if not os.path.exists(output_dir): | |
os.makedirs(output_dir) | |
images.save(os.path.join(output_dir, f"{current_date}-{seed}.jpg")) | |
return os.path.join(output_dir, f"{current_date}-{seed}.jpg") | |
iface = gr.Interface( | |
fn=process, | |
inputs=[ | |
gr.Image(label="Upload Image"), | |
gr.Textbox(label="prompt",placeholder="A man, in a forest, adventuring"), | |
gr.Textbox(label="negative prompt",placeholder="monochrome, lowres, bad anatomy, worst quality, low quality, blurry"), | |
], | |
outputs=[ | |
gr.Image(label="Output"), | |
], | |
title="ConsistentID Demo", | |
description="Put reference portrait below" | |
) | |
iface.launch() # zero.device | |
# @spaces.GPU | |
# def greet(n): | |
# print(zero.device) # <-- 'cuda:0' π€ | |
# return f"Hello {zero + n} Tensor" | |
# demo = gr.Interface(fn=greet, inputs=gr.Number(), outputs=gr.Text()) | |
# demo.launch() | |