atifnadeem commited on
Commit
33a154e
·
verified ·
1 Parent(s): 50b1574

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +194 -0
app.py ADDED
@@ -0,0 +1,194 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import random
3
+ import uuid
4
+ from typing import Tuple
5
+ import gradio as gr
6
+ import numpy as np
7
+ from PIL import Image
8
+ import spaces
9
+ import torch
10
+ from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
11
+ from huggingface_hub import login
12
+ # Log in to Hugging Face using the provided token
13
+ hf_token = ''
14
+ login(hf_token)
15
+
16
+ DESCRIPTIONz = """## STABLE IMAGINE 🍺"""
17
+
18
+ def save_image(img):
19
+ unique_name = str(uuid.uuid4()) + ".png"
20
+ img.save(unique_name)
21
+ return unique_name
22
+
23
+ def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
24
+ if randomize_seed:
25
+ seed = random.randint(0, MAX_SEED)
26
+ return seed
27
+
28
+ MAX_SEED = np.iinfo(np.int32).max
29
+ DESCRIPTIONz = ""
30
+
31
+ if not torch.cuda.is_available():
32
+ DESCRIPTIONz += """
33
+ <p>⚠️Running on CPU, This may not work on CPU. If it runs for an extended time or if you encounter errors, try running it on a GPU by duplicating the space using @spaces.GPU(). 📍</p>
34
+ """
35
+
36
+ USE_TORCH_COMPILE = 0
37
+ ENABLE_CPU_OFFLOAD = 0
38
+
39
+ if torch.cuda.is_available():
40
+ pipe = StableDiffusionXLPipeline.from_pretrained(
41
+ "SG161222/RealVisXL_V4.0_Lightning",
42
+ torch_dtype=torch.float16,
43
+ use_safetensors=True,
44
+ )
45
+ pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
46
+ pipe.to("cuda")
47
+
48
+ if USE_TORCH_COMPILE:
49
+ pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
50
+
51
+ else:
52
+ # If CUDA is not available, fall back to CPU (not ideal for SDXL)
53
+ pipe = StableDiffusionXLPipeline.from_pretrained(
54
+ "SG161222/RealVisXL_V4.0_Lightning",
55
+ torch_dtype=torch.float32, # safer for CPU
56
+ use_safetensors=True,
57
+ )
58
+ pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
59
+ pipe.to("cpu")
60
+
61
+ if ENABLE_CPU_OFFLOAD:
62
+ # Optionally offload to CPU with accelerate or similar, if set up
63
+ pipe.enable_model_cpu_offload()
64
+
65
+
66
+
67
+ LORA_OPTIONS = {
68
+ "Realism (face/character)👦🏻": ("prithivMLmods/Canopus-Realism-LoRA", "Canopus-Realism-LoRA.safetensors", "rlms"),
69
+ "Pixar (art/toons)🙀": ("prithivMLmods/Canopus-Pixar-Art", "Canopus-Pixar-Art.safetensors", "pixar"),
70
+ "Interior Architecture (house/hotel)🏠": ("prithivMLmods/Canopus-Interior-Architecture-0.1", "Canopus-Interior-Architecture-0.1δ.safetensors", "arch"),
71
+ "Fashion Product (wearing/usable)👜": ("prithivMLmods/Canopus-Fashion-Product-Dilation", "Canopus-Fashion-Product-Dilation.safetensors", "fashion"),
72
+ "Minimalistic Image (minimal/detailed)🏞️": ("prithivMLmods/Pegasi-Minimalist-Image-Style", "Pegasi-Minimalist-Image-Style.safetensors", "minimalist"),
73
+ "Modern Clothing (trend/new)👕": ("prithivMLmods/Canopus-Modern-Clothing-Design", "Canopus-Modern-Clothing-Design.safetensors", "mdrnclth"),
74
+ "Animaliea (farm/wild)🫎": ("prithivMLmods/Canopus-Animaliea-Artism", "Canopus-Animaliea-Artism.safetensors", "Animaliea"),
75
+ "Canes Cars (realistic/futurecars)🚘": ("prithivMLmods/Canes-Cars-Model-LoRA", "Canes-Cars-Model-LoRA.safetensors", "car"),
76
+ "Art Minimalistic (paint/semireal)🎨": ("prithivMLmods/Canopus-Art-Medium-LoRA", "Canopus-Art-Medium-LoRA.safetensors", "mdm"),
77
+ }
78
+
79
+ for model_name, weight_name, adapter_name in LORA_OPTIONS.values():
80
+ pipe.load_lora_weights(model_name, weight_name=weight_name, adapter_name=adapter_name)
81
+ pipe.to("cuda")
82
+
83
+
84
+ style_list = [
85
+ {
86
+ "name": "3840 x 2160",
87
+ "prompt": "hyper-realistic 8K image of {prompt}. ultra-detailed, lifelike, high-resolution, sharp, vibrant colors, photorealistic",
88
+ "negative_prompt": "cartoonish, low resolution, blurry, simplistic, abstract, deformed, ugly",
89
+ },
90
+ # Add more style dicts here if needed
91
+ ]
92
+
93
+ styles = {k["name"]: (k["prompt"], k["negative_prompt"]) for k in style_list}
94
+
95
+ def apply_style(style_name: str, positive: str, negative: str = "") -> Tuple[str, str]:
96
+ if style_name in styles:
97
+ p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
98
+ else:
99
+ p, n = styles[DEFAULT_STYLE_NAME]
100
+
101
+ if not negative:
102
+ negative = ""
103
+ return p.replace("{prompt}", positive), n + negative
104
+
105
+
106
+ DEFAULT_STYLE_NAME = "3840 x 2160"
107
+
108
+
109
+ @spaces.GPU(duration=60, enable_queue=True)
110
+ def generate(
111
+ prompt: str,
112
+ negative_prompt: str = "",
113
+ use_negative_prompt: bool = False,
114
+ seed: int = 0,
115
+ width: int = 1024,
116
+ height: int = 1024,
117
+ guidance_scale: float = 3,
118
+ randomize_seed: bool = False,
119
+ style_name: str = DEFAULT_STYLE_NAME,
120
+ lora_model: str = "Realism (face/character)👦🏻",
121
+ progress=gr.Progress(track_tqdm=True),
122
+ ):
123
+ seed = int(randomize_seed_fn(seed, randomize_seed))
124
+
125
+ positive_prompt, effective_negative_prompt = apply_style(style_name, prompt, negative_prompt)
126
+
127
+ if not use_negative_prompt:
128
+ effective_negative_prompt = "" # type: ignore
129
+
130
+ model_name, weight_name, adapter_name = LORA_OPTIONS[lora_model]
131
+ pipe.set_adapters(adapter_name)
132
+
133
+ images = pipe(
134
+ prompt=positive_prompt,
135
+ negative_prompt=effective_negative_prompt,
136
+ width=width,
137
+ height=height,
138
+ guidance_scale=guidance_scale,
139
+ num_inference_steps=20,
140
+ num_images_per_prompt=1,
141
+ cross_attention_kwargs={"scale": 0.65},
142
+ output_type="pil",
143
+ ).images
144
+ image_paths = [save_image(img) for img in images]
145
+ return image_paths, seed
146
+
147
+
148
+ with gr.Blocks() as demo:
149
+ gr.Markdown(DESCRIPTIONz)
150
+
151
+ with gr.Row():
152
+ input_prompt = gr.Textbox(label="Prompt", placeholder="Enter prompt", lines=2)
153
+ use_negative_prompt = gr.Checkbox(label="Use negative prompt?", value=False)
154
+ negative_prompt = gr.Textbox(label="Negative Prompt", placeholder="Enter negative prompt", lines=2)
155
+
156
+ with gr.Row():
157
+ randomize_seed = gr.Checkbox(label="Randomize Seed", value=False)
158
+ seed = gr.Number(value=0, label="Seed")
159
+
160
+ with gr.Row():
161
+ style_dropdown = gr.Dropdown(label="Image Style", choices=list(styles.keys()), value=DEFAULT_STYLE_NAME)
162
+ lora_dropdown = gr.Dropdown(label="LoRA Model", choices=list(LORA_OPTIONS.keys()), value="Realism (face/character)👦🏻")
163
+
164
+ with gr.Row():
165
+ width = gr.Slider(512, 2048, value=1024, step=64, label="Width")
166
+ height = gr.Slider(512, 2048, value=1024, step=64, label="Height")
167
+
168
+ with gr.Row():
169
+ guidance_scale = gr.Slider(1.0, 15.0, value=3, step=0.5, label="Guidance Scale")
170
+
171
+ output_gallery = gr.Gallery(label="Generated Images", columns=[2], height="auto")
172
+ output_seed = gr.Number(label="Final Seed", interactive=False)
173
+
174
+
175
+ generate_button = gr.Button("Generate Images")
176
+
177
+ generate_button.click(
178
+ fn=generate,
179
+ inputs=[
180
+ input_prompt,
181
+ negative_prompt,
182
+ use_negative_prompt,
183
+ seed,
184
+ width,
185
+ height,
186
+ guidance_scale,
187
+ randomize_seed,
188
+ style_dropdown,
189
+ lora_dropdown,
190
+ ],
191
+ outputs=[output_gallery, output_seed],
192
+ )
193
+
194
+ demo.launch()