|
|
|
|
|
import os |
|
import random |
|
import uuid |
|
import base64 |
|
import numpy as np |
|
from PIL import Image |
|
import torch |
|
import streamlit as st |
|
|
|
from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler |
|
|
|
DESCRIPTION = """# DALL•E 3 XL v2 High Fi""" |
|
|
|
def create_download_link(filename): |
|
with open(filename, "rb") as file: |
|
encoded_string = base64.b64encode(file.read()).decode('utf-8') |
|
download_link = f'<a href="data:image/png;base64,{encoded_string}" download="{filename}">Download Image</a>' |
|
return download_link |
|
|
|
def save_image(img, prompt): |
|
unique_name = str(uuid.uuid4()) + ".png" |
|
img.save(unique_name) |
|
|
|
|
|
filename = f"{prompt}.png" |
|
img.save(filename) |
|
return filename |
|
|
|
def randomize_seed_fn(seed: int, randomize_seed: bool) -> int: |
|
if randomize_seed: |
|
seed = random.randint(0, MAX_SEED) |
|
return seed |
|
|
|
MAX_SEED = np.iinfo(np.int32).max |
|
|
|
if not torch.cuda.is_available(): |
|
DESCRIPTION += "\n<p>Running on CPU 🥶 This demo may not work on CPU.</p>" |
|
|
|
MAX_SEED = np.iinfo(np.int32).max |
|
|
|
USE_TORCH_COMPILE = 0 |
|
ENABLE_CPU_OFFLOAD = 0 |
|
|
|
|
|
|
|
if torch.cuda.is_available(): |
|
pipe = StableDiffusionXLPipeline.from_pretrained( |
|
"fluently/Fluently-XL-v4", |
|
torch_dtype=torch.float16, |
|
use_safetensors=True, |
|
) |
|
pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config) |
|
|
|
|
|
pipe.load_lora_weights("ehristoforu/dalle-3-xl-v2", weight_name="dalle-3-xl-lora-v2.safetensors", adapter_name="dalle") |
|
pipe.set_adapters("dalle") |
|
|
|
pipe.to("cuda") |
|
|
|
def generate( |
|
prompt: str, |
|
negative_prompt: str = "", |
|
use_negative_prompt: bool = False, |
|
seed: int = 0, |
|
width: int = 1024, |
|
height: int = 1024, |
|
guidance_scale: float = 3, |
|
randomize_seed: bool = False, |
|
): |
|
seed = int(randomize_seed_fn(seed, randomize_seed)) |
|
|
|
if not use_negative_prompt: |
|
negative_prompt = "" |
|
|
|
images = pipe( |
|
prompt=prompt, |
|
negative_prompt=negative_prompt, |
|
width=width, |
|
height=height, |
|
guidance_scale=guidance_scale, |
|
num_inference_steps=20, |
|
num_images_per_prompt=1, |
|
cross_attention_kwargs={"scale": 0.65}, |
|
output_type="pil", |
|
).images |
|
image_paths = [save_image(img, prompt) for img in images] |
|
|
|
download_links = [create_download_link(path) for path in image_paths] |
|
|
|
print(image_paths) |
|
return image_paths, seed, download_links |
|
|
|
if run_button: |
|
image_paths, seed, download_links = generate( |
|
prompt=prompt, |
|
negative_prompt=negative_prompt, |
|
use_negative_prompt=use_negative_prompt, |
|
seed=seed, |
|
width=width, |
|
height=height, |
|
guidance_scale=guidance_scale, |
|
randomize_seed=randomize_seed, |
|
) |
|
|
|
for image_path in image_paths: |
|
st.image(image_path, caption=prompt) |
|
|
|
for download_link in download_links: |
|
st.markdown(download_link, unsafe_allow_html=True) |
|
|
|
st.text(f"Seed: {seed}") |
|
else: |
|
st.warning("CUDA is not available. The demo may not work on CPU.") |
|
|
|
examples = [ |
|
"a modern hospital room with advanced medical equipment and a patient resting comfortably", |
|
"a team of surgeons performing a delicate operation using state-of-the-art surgical robots", |
|
"a elderly woman smiling while a nurse checks her vital signs using a holographic display", |
|
"a child receiving a painless vaccination from a friendly robot nurse in a colorful pediatric clinic", |
|
"a group of researchers working in a high-tech laboratory, developing new treatments for rare diseases", |
|
"a telemedicine consultation between a doctor and a patient, using virtual reality technology for a immersive experience" |
|
] |
|
|
|
|
|
st.markdown(DESCRIPTION) |
|
|
|
with st.form(key="generation_form"): |
|
prompt = st.text_input("Prompt", max_chars=None, placeholder="Enter your prompt") |
|
use_negative_prompt = st.checkbox("Use negative prompt", value=True) |
|
if use_negative_prompt: |
|
negative_prompt = st.text_area( |
|
"Negative prompt", |
|
value="""(deformed, distorted, disfigured:1.3), poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, (mutated hands and fingers:1.4), disconnected limbs, mutation, mutated, ugly, disgusting, blurry, amputation, (NSFW:1.25)""", |
|
placeholder="Enter a negative prompt", |
|
height=100, |
|
) |
|
else: |
|
negative_prompt = "" |
|
|
|
col1, col2 = st.columns(2) |
|
with col1: |
|
width = st.slider("Width", min_value=512, max_value=2048, step=8, value=1920) |
|
with col2: |
|
height = st.slider("Height", min_value=512, max_value=2048, step=8, value=1080) |
|
|
|
col3, col4 = st.columns(2) |
|
with col3: |
|
guidance_scale = st.slider("Guidance Scale", min_value=0.1, max_value=20.0, step=0.1, value=20.0) |
|
with col4: |
|
randomize_seed = st.checkbox("Randomize seed", value=True) |
|
if not randomize_seed: |
|
seed = st.slider("Seed", min_value=0, max_value=MAX_SEED, step=1, value=0) |
|
else: |
|
seed = 0 |
|
|
|
run_button = st.form_submit_button("Run") |
|
|
|
def set_prompt(example): |
|
prompt.value = example |
|
|
|
st.subheader("Examples") |
|
for example in examples: |
|
st.button(example, key=example, on_click=set_prompt, args=(example,)) |