File size: 10,639 Bytes
bb1b13b
 
8783961
bb1b13b
de9d198
bb1b13b
 
 
 
 
 
02659a8
3d80262
25d3a18
e3d2786
bb1b13b
de9d198
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
02659a8
 
 
 
 
 
6a7aae2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
02659a8
 
 
 
 
 
 
 
6fb452a
3d80262
e67a76c
 
 
 
 
45337d4
e67a76c
 
 
6d1124c
e67a76c
fbab0f5
 
 
 
 
 
 
 
e67a76c
793ed3f
 
 
fbab0f5
 
 
793ed3f
 
02659a8
 
fbab0f5
6d1124c
dffd00a
fbab0f5
de9b7a6
02659a8
 
 
 
de9d198
 
 
 
 
 
 
 
 
 
bd88c3e
de9d198
bd88c3e
de9d198
 
 
 
bd88c3e
de9d198
 
 
 
 
 
 
 
 
 
 
bd88c3e
 
de9d198
ed98223
 
 
ae4aefb
ed98223
 
 
 
 
e6d5b6f
ed98223
3043030
 
 
 
 
 
ed98223
bd88c3e
de9d198
bd88c3e
de9d198
bd88c3e
de9d198
 
 
bd88c3e
 
de9d198
 
 
bd88c3e
de9d198
 
 
 
 
 
feb562f
c89f94f
329ec73
 
 
 
 
 
 
 
 
 
 
 
5187dd9
 
 
 
 
329ec73
 
8678edc
 
de9d198
 
 
 
 
5366491
de9d198
 
12fd800
de9d198
 
 
 
 
a8114ef
de9d198
 
eb403fc
de9d198
 
 
 
 
5366491
de9d198
 
bd88c3e
3043030
80e4491
3043030
ec71404
3043030
 
 
 
 
 
 
ec71404
3043030
ec71404
6242eec
ec71404
fec2ec7
 
 
 
 
 
ec71404
de9d198
 
 
 
 
 
a1c7876
de9d198
ec71404
 
e9f4989
bd88c3e
 
8f40af2
 
 
 
80e4491
ba79244
bd88c3e
80e4491
bd88c3e
80e4491
cfdc849
80e4491
5e7f076
80e4491
 
 
bd88c3e
 
 
ec71404
de9d198
bd88c3e
3043030
bd88c3e
 
 
 
de9d198
d0b4c7a
 
ef5e5f7
 
 
 
 
d0b4c7a
 
de9d198
 
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
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
import base64
import datetime
import gradio as gr
import numpy as np
import os
import pytz
import psutil
import re
import random
import torch
import time
import shutil  # Added for zip functionality
import zipfile
from PIL import Image
from io import BytesIO
from diffusers import DiffusionPipeline, LCMScheduler, AutoencoderTiny

try:
    import intel_extension_for_pytorch as ipex
except:
    pass

SAFETY_CHECKER = os.environ.get("SAFETY_CHECKER", None)
TORCH_COMPILE = os.environ.get("TORCH_COMPILE", None)
HF_TOKEN = os.environ.get("HF_TOKEN", None)
# check if MPS is available OSX only M1/M2/M3 chips
mps_available = hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
xpu_available = hasattr(torch, "xpu") and torch.xpu.is_available()
device = torch.device(
    "cuda" if torch.cuda.is_available() else "xpu" if xpu_available else "cpu"
)
torch_device = device
torch_dtype = torch.float16

# Function to encode a file to base64
def encode_file_to_base64(file_path):
    with open(file_path, "rb") as file:
        encoded = base64.b64encode(file.read()).decode()
    return encoded

def create_zip_of_files(files):
    """
    Create a zip file from a list of files.
    """
    zip_name = "all_files.zip"
    with zipfile.ZipFile(zip_name, 'w') as zipf:
        for file in files:
            zipf.write(file)
    return zip_name


def get_zip_download_link(zip_file):
    """
    Generate a link to download the zip file.
    """
    with open(zip_file, 'rb') as f:
        data = f.read()
    b64 = base64.b64encode(data).decode()
    href = f'<a href="data:application/zip;base64,{b64}" download="{zip_file}">Download All</a>'
    return href

# Function to clear all image files
def clear_all_images():
    base_dir = os.getcwd()  # Get the current base directory
    img_files = [file for file in os.listdir(base_dir) if file.lower().endswith((".png", ".jpg", ".jpeg"))]  # List all files ending with ".jpg" or ".jpeg"
    
    # Remove all image files
    for file in img_files:
        os.remove(file)
        print('removed:' + file)
  
# add file save and download and clear:
# Function to create a zip file from a list of files
def create_zip(files):
    timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
    zip_filename = f"images_{timestamp}.zip"
    print('Creating file ' + zip_filename)
    with zipfile.ZipFile(zip_filename, 'w') as zipf:
        for file in files:
            zipf.write(file, os.path.basename(file))
            print('added:' + file)
    return zip_filename

def get_zip_download_link(zip_file):
    """
    Generate a link to download the zip file.
    """
    zip_base64 = encode_file_to_base64(zip_file)  # Encode the zip file to base64
    href = f'<a href="data:application/zip;base64,{zip_base64}" download="{zip_file}">Download All</a>'
    return href
    
def save_all_images(images):
    if len(images) == 0:
        return None, None
    zip_filename = create_zip_of_files(images)  # Create a zip file from the list of image files
    print(f"Zip file created: {zip_filename}")
    download_link = get_zip_download_link(zip_filename)
    return zip_filename, download_link
        
def save_all_button_click():
    images = [file for file in os.listdir() if file.lower().endswith((".png", ".jpg", ".jpeg"))]
    zip_filename, download_link = save_all_images(images)
    if download_link:
        gr.HTML(download_link)


# Function to handle "Clear All" button click
def clear_all_button_click():
    clear_all_images()

print(f"SAFETY_CHECKER: {SAFETY_CHECKER}")
print(f"TORCH_COMPILE: {TORCH_COMPILE}")
print(f"device: {device}")

if mps_available:
    device = torch.device("mps")
    torch_device = "cpu"
    torch_dtype = torch.float32

if SAFETY_CHECKER == "True":
    pipe = DiffusionPipeline.from_pretrained("Lykon/dreamshaper-7")
else:
    pipe = DiffusionPipeline.from_pretrained("Lykon/dreamshaper-7", safety_checker=None)

pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
pipe.to(device=torch_device, dtype=torch_dtype).to(device)
pipe.unet.to(memory_format=torch.channels_last)
pipe.set_progress_bar_config(disable=True)

# check if computer has less than 64GB of RAM using sys or os
if psutil.virtual_memory().total < 64 * 1024**3:
    pipe.enable_attention_slicing()

if TORCH_COMPILE:
    pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
    pipe.vae = torch.compile(pipe.vae, mode="reduce-overhead", fullgraph=True)
    pipe(prompt="warmup", num_inference_steps=1, guidance_scale=8.0)

# Load LCM LoRA
pipe.load_lora_weights("latent-consistency/lcm-lora-sdv1-5")
pipe.fuse_lora()

def safe_filename(text):
    """Generate a safe filename from a string."""
    safe_text = re.sub(r'\W+', '_', text)
    timestamp = datetime.datetime.now().strftime("%Y%m%d")
    return f"{safe_text}_{timestamp}.png"
    
def encode_image(image):
    """Encode image to base64."""
    buffered = BytesIO()
    #image.save(buffered, format="PNG")
    return base64.b64encode(buffered.getvalue()).decode()

def fake_gan():
    base_dir = os.getcwd()  # Get the current base directory
    img_files = [file for file in os.listdir(base_dir) if file.lower().endswith((".png", ".jpg", ".jpeg"))]  # List all files ending with ".jpg" or ".jpeg"
    images = [(random.choice(img_files), os.path.splitext(file)[0]) for file in img_files]
    return images
    
def predict(prompt, guidance, steps, seed=1231231):
    generator = torch.manual_seed(seed)
    last_time = time.time()
    results = pipe(
        prompt=prompt,
        generator=generator,
        num_inference_steps=steps,
        guidance_scale=guidance,
        width=512,
        height=512,
        # original_inference_steps=params.lcm_steps,
        output_type="pil",
    )
    print(f"Pipe took {time.time() - last_time} seconds")
    nsfw_content_detected = (
        results.nsfw_content_detected[0]
        if "nsfw_content_detected" in results
        else False
    )
    if nsfw_content_detected:
        nsfw=gr.Button("🕹️NSFW🎨", scale=1)

    try: 
        central = pytz.timezone('US/Central')
        safe_date_time = datetime.datetime.now().strftime("%Y%m%d")
        replaced_prompt = prompt.replace(" ", "_").replace("\n", "_")
        safe_prompt = "".join(x for x in replaced_prompt if x.isalnum() or x == "_")[:90]
        filename = f"{safe_date_time}_{safe_prompt}.png"
        
        # Save the image
        if len(results.images) > 0:
            image_path = os.path.join("", filename)  # Specify your directory
            results.images[0].save(image_path)
            print(f"#Image saved as {image_path}")
            gr.File(image_path)
            gr.Button(link=image_path)
            # encoded_image = encode_image(image)
            # html_link = f'<a href="data:image/png;base64,{encoded_image}" download="{filename}">Download Image</a>'
            # gr.HTML(html_link)
    except:
        return results.images[0]

    return results.images[0] if len(results.images) > 0 else None


css = """
#container{
    margin: 0 auto;
    max-width: 40rem;
}
#intro{
    max-width: 100%;
    text-align: center;
    margin: 0 auto;
}
"""
with gr.Blocks(css=css) as demo:

    with gr.Column(elem_id="container"):
        gr.Markdown(
            """4📝RT🖼️Images - 🕹️ Real Time 🎨 Image Generator Gallery 🌐""",
            elem_id="intro",
        )
        with gr.Row():
            with gr.Row():
                prompt = gr.Textbox(
                    placeholder="Insert your prompt here:", scale=5, container=False
                )
                generate_bt = gr.Button("Generate", scale=1)

        # Image Result from last prompt
        image = gr.Image(type="filepath")

        # Gallery of Generated Images with Image Names in Random Set to Download
        with gr.Row(variant="compact"):
            text = gr.Textbox(
                label="Image Sets",
                show_label=False,
                max_lines=1,
                placeholder="Enter your prompt",
            )
            btn = gr.Button("Generate Gallery of Saved Images")
        gallery = gr.Gallery(
            label="Generated Images", show_label=False, elem_id="gallery"
        )

        with gr.Row(variant="compact"):
            # Add "Save All" button with emoji
            save_all_button = gr.Button("💾 Save All", scale=1)
            # Add "Clear All" button with emoji
            clear_all_button = gr.Button("🗑️ Clear All", scale=1)

        # Advanced Generate Options
        with gr.Accordion("Advanced options", open=False):
            guidance = gr.Slider(
                label="Guidance", minimum=0.0, maximum=5, value=0.3, step=0.001
            )
            steps = gr.Slider(label="Steps", value=4, minimum=2, maximum=10, step=1)
            seed = gr.Slider(
                randomize=True, minimum=0, maximum=12013012031030, label="Seed", step=1
            )

        # Diffusers
        with gr.Accordion("Run with diffusers"):
            gr.Markdown(
                """## Running LCM-LoRAs it with `diffusers`
            ```bash
            pip install diffusers==0.23.0
            ```
            
            ```py
            from diffusers import DiffusionPipeline, LCMScheduler
            pipe = DiffusionPipeline.from_pretrained("Lykon/dreamshaper-7").to("cuda") 
            pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
            pipe.load_lora_weights("latent-consistency/lcm-lora-sdv1-5") #yes, it's a normal LoRA
            results = pipe(
                prompt="ImageEditor",
                num_inference_steps=4,
                guidance_scale=0.0,
            )
            results.images[0]
            ```
            """
            )

        # Function IO Eventing and Controls
        inputs = [prompt, guidance, steps, seed]
        generate_bt.click(fn=predict, inputs=inputs, outputs=image, show_progress=False)
        btn.click(fake_gan, None, gallery)
        prompt.input(fn=predict, inputs=inputs, outputs=image, show_progress=False)
        guidance.change(fn=predict, inputs=inputs, outputs=image, show_progress=False)
        steps.change(fn=predict, inputs=inputs, outputs=image, show_progress=False)
        seed.change(fn=predict, inputs=inputs, outputs=image, show_progress=False)

        # Attach click event handlers to the buttons
        save_all_button.click(save_all_button_click)

        with gr.Column():
            file_obj = gr.File(label="Input File")
            input= file_obj
                    
        clear_all_button.click(clear_all_button_click)

demo.queue()
demo.launch()