File size: 9,158 Bytes
f6b9815
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
66f8fc1
c392e6f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
7fa88e5
 
c392e6f
f6b9815
 
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
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
run_api = False
SSD_1B = False
import os

# Use GPU
gpu_info = os.popen("nvidia-smi").read()
if "failed" in gpu_info:
    print("Not connected to a GPU")
    is_gpu = False
else:
    print(gpu_info)
    is_gpu = True
print(is_gpu)


from IPython.display import clear_output


def check_enviroment():
    try:
        import torch

        print("Enviroment is already installed.")
    except ImportError:
        print("Enviroment not found. Installing...")
        # Install requirements from requirements.txt
        os.system("pip install -r requirements.txt")
        # Install gradio version 3.48.0
        os.system("pip install gradio==3.39.0")
        # Install python-dotenv
        os.system("pip install python-dotenv")
        # Clear the output
        clear_output()

        print("Enviroment installed successfully.")


# Call the function to check and install Packages if necessary
check_enviroment()


from IPython.display import clear_output
import os
import gradio as gr
import numpy as np
import PIL
import base64
import io
import torch
from diffusers import UNet2DConditionModel, DiffusionPipeline, LCMScheduler

# SDXL
from diffusers import UNet2DConditionModel, DiffusionPipeline, LCMScheduler

# Get the current directory
current_dir = os.getcwd()
model_path = os.path.join(current_dir)
# Set the cache path
cache_path = os.path.join(current_dir, "cache")
MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "1024"))
SECRET_TOKEN = os.getenv("SECRET_TOKEN", "default_secret")

# Uncomment the following line if you are using PyTorch 1.10 or later
# os.environ["TORCH_USE_CUDA_DSA"] = "1"

if is_gpu:
    # Uncomment the following line if you want to enable CUDA launch blocking
    os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
else:
    # Uncomment the following line if you want to use CPU instead of GPU
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")


# Get the current directory
current_dir = os.getcwd()
model_path = os.path.join(current_dir)

# Set the cache path
cache_path = os.path.join(current_dir, "cache")

if not SSD_1B:

    unet = UNet2DConditionModel.from_pretrained(
        "latent-consistency/lcm-sdxl",
        torch_dtype=torch.float16,
        variant="fp16",
        cache_dir=cache_path,
    )
    pipe = DiffusionPipeline.from_pretrained(
        # "stabilityai/stable-diffusion-xl-base-1.0",
        "stabilityai/sdxl-turbo",
        unet=unet,
        torch_dtype=torch.float16,
        variant="fp16",
        cache_dir=cache_path,
    )

    pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
    if torch.cuda.is_available():
        pipe.to("cuda")
else:
    # SSD-1B
    from diffusers import LCMScheduler, AutoPipelineForText2Image

    pipe = AutoPipelineForText2Image.from_pretrained(
        "segmind/SSD-1B",
        torch_dtype=torch.float16,
        variant="fp16",
        cache_dir=cache_path,
    )
    pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
    if torch.cuda.is_available():
        pipe.to("cuda")

    # load and fuse
    pipe.load_lora_weights("latent-consistency/lcm-lora-ssd-1b")
    pipe.fuse_lora()


def generate(
    prompt: str,
    negative_prompt: str = "",
    seed: int = 0,
    width: int = 1024,
    height: int = 1024,
    guidance_scale: float = 0.0,
    num_inference_steps: int = 4,
    secret_token: str = "",
) -> PIL.Image.Image:
    if secret_token != SECRET_TOKEN:
        raise gr.Error(
            f"Invalid secret token. Please fork the original space if you want to use it for yourself."
        )

    generator = torch.Generator().manual_seed(seed)

    image = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        width=width,
        height=height,
        guidance_scale=guidance_scale,
        num_inference_steps=num_inference_steps,
        generator=generator,
        output_type="pil",
    ).images[0]
    return image


clear_output()

from IPython.display import display


def generate_image(prompt="A beautiful and sexy girl"):
    # Generate the image using the prompt
    generated_image = generate(
        prompt=prompt,
        negative_prompt="",
        seed=0,
        width=1024,
        height=1024,
        guidance_scale=0.0,
        num_inference_steps=4,
        secret_token="default_secret",  # Replace with your secret token
    )
    # Display the image in the Jupyter Notebook
    display(generated_image)


if not run_api:
    secret_token = gr.Text(
        label="Secret Token",
        max_lines=1,
        placeholder="Enter your secret token",
    )
    prompt = gr.Text(
        label="Prompt",
        show_label=False,
        max_lines=1,
        placeholder="Enter your prompt",
        container=False,
    )
    result = gr.Image(label="Result", show_label=False)
    negative_prompt = gr.Text(
        label="Negative prompt",
        max_lines=1,
        placeholder="Enter a negative prompt",
        visible=True,
    )
    seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)

    width = gr.Slider(
        label="Width",
        minimum=256,
        maximum=MAX_IMAGE_SIZE,
        step=32,
        value=1024,
    )
    height = gr.Slider(
        label="Height",
        minimum=256,
        maximum=MAX_IMAGE_SIZE,
        step=32,
        value=1024,
    )
    guidance_scale = gr.Slider(
        label="Guidance scale", minimum=0, maximum=2, step=0.1, value=0.0
    )
    num_inference_steps = gr.Slider(
        label="Number of inference steps", minimum=1, maximum=8, step=1, value=4
    )
    inputs = [
        prompt,
        negative_prompt,
        seed,
        width,
        height,
        guidance_scale,
        num_inference_steps,
        secret_token,
    ]
    iface = gr.Interface(
        fn=generate,
        inputs=inputs,
        outputs=result,
        title="Image Generator",
        description="Generate images based on prompts.",
    )

    iface.launch()


if run_api:
    with gr.Blocks() as demo:
        gr.HTML(
            """
        <div style="z-index: 100; position: fixed; top: 0px; right: 0px; left: 0px; bottom: 0px; width: 100%; height: 100%; background: white; display: flex; align-items: center; justify-content: center; color: black;">
            <div style="text-align: center; color: black;">
                <p style="color: black;">This space is a REST API to programmatically generate images using LCM LoRA SSD-1B.</p>
                <p style="color: black;">It is not meant to be directly used through a user interface, but using code and an access key.</p>
            </div>
        </div>"""
        )
        secret_token = gr.Text(
            label="Secret Token",
            max_lines=1,
            placeholder="Enter your secret token",
        )
        prompt = gr.Text(
            label="Prompt",
            show_label=False,
            max_lines=1,
            placeholder="Enter your prompt",
            container=False,
        )
        result = gr.Image(label="Result", show_label=False)
        negative_prompt = gr.Text(
            label="Negative prompt",
            max_lines=1,
            placeholder="Enter a negative prompt",
            visible=True,
        )
        seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)

        width = gr.Slider(
            label="Width",
            minimum=256,
            maximum=MAX_IMAGE_SIZE,
            step=32,
            value=1024,
        )
        height = gr.Slider(
            label="Height",
            minimum=256,
            maximum=MAX_IMAGE_SIZE,
            step=32,
            value=1024,
        )
        guidance_scale = gr.Slider(
            label="Guidance scale", minimum=0, maximum=2, step=0.1, value=0.0
        )
        num_inference_steps = gr.Slider(
            label="Number of inference steps", minimum=1, maximum=8, step=1, value=4
        )

        inputs = [
            prompt,
            negative_prompt,
            seed,
            width,
            height,
            guidance_scale,
            num_inference_steps,
            secret_token,
        ]
        prompt.submit(
            fn=generate,
            inputs=inputs,
            outputs=result,
            api_name="run",
        )

    # demo.queue(max_size=32).launch()
    # Launch the Gradio app with multiple workers and debug mode enabled
    # demo.queue(max_size=32).launch(debug=True)# For Standard
    demo.queue(max_size=32).launch(server_name="0.0.0.0", server_port=7860)  # Docker


'''


import gradio as gr
import subprocess
def run_command(command):
    try:
        result = subprocess.check_output(command, shell=True, text=True)
        return result
    except subprocess.CalledProcessError as e:
        return f"Error: {e}"
iface = gr.Interface(
    fn=run_command,
    inputs="text",
    outputs="text",
    #live=True,
    title="Command Output Viewer",
    description="Enter a command and view its output.",
    examples=[
    ["ls"],
    ["pwd"],
    ["echo 'Hello, Gradio!'"],
    ["python --version"]]
)
iface.launch(server_name="0.0.0.0", server_port=7860)
'''