marahmerah commited on
Commit
20086a8
·
verified ·
1 Parent(s): 1b9717a

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +70 -58
app.py CHANGED
@@ -2,44 +2,57 @@ import gradio as gr
2
  import torch
3
  import random
4
  import os
5
- import time
6
  from PIL import Image
7
  from deep_translator import GoogleTranslator
8
  from diffusers import DiffusionPipeline
9
- from huggingface_hub import hf_hub_download
10
 
11
- # Project by Nymbo with LoRA integration
 
 
12
 
13
- # Model and LoRA configuration
14
  BASE_MODEL = "black-forest-labs/FLUX.1-dev"
15
  LORA_REPO = "burhansyam/uncen"
16
- LORA_WEIGHTS_NAME = "uncen.safetensors" # Adjust if different
17
- torch_dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
18
 
19
- # Initialize the pipeline with LoRA
20
  def init_pipeline():
 
21
  pipe = DiffusionPipeline.from_pretrained(
22
  BASE_MODEL,
23
- torch_dtype=torch_dtype
 
24
  )
25
 
26
- # Load LoRA weights
27
- pipe.load_lora_weights(
28
- hf_hub_download(repo_id=LORA_REPO, filename=LORA_WEIGHTS_NAME),
29
- adapter_name="uncen"
 
30
  )
 
31
 
32
- # Enable model offloading if needed
33
  if torch.cuda.is_available():
34
  pipe.to("cuda")
35
- pipe.enable_xformers_memory_efficient_attention()
 
 
 
36
 
37
  return pipe
38
 
39
- pipe = init_pipeline()
 
 
 
 
40
 
41
  def convert_to_png(image):
42
- """Convert any image format to true PNG format"""
43
  png_buffer = io.BytesIO()
44
  if image.mode == 'RGBA':
45
  image.save(png_buffer, format='PNG', optimize=True)
@@ -50,64 +63,64 @@ def convert_to_png(image):
50
  png_buffer.seek(0)
51
  return Image.open(png_buffer)
52
 
53
- def query(prompt, is_negative=False, steps=35, cfg_scale=7, sampler="DPM++ 2M Karras",
54
- seed=-1, strength=0.7, width=1024, height=1024):
 
 
 
 
 
 
 
 
55
  if not prompt:
56
  return None
57
 
58
- key = random.randint(0, 999)
59
-
60
- # Translate prompt
61
  try:
62
- prompt = GoogleTranslator(source='id', target='en').translate(prompt)
63
- print(f'\033[1mGeneration {key} translation:\033[0m {prompt}')
64
- prompt = f"{prompt} | ultra detail, ultra elaboration, ultra quality, perfect."
65
- except Exception as e:
66
- print(f"Translation error: {e}")
67
 
68
- print(f'\033[1mGeneration {key}:\033[0m {prompt}')
69
-
70
- # Set random seed if not specified
71
- generator = None
72
- if seed != -1:
73
- generator = torch.Generator(device="cuda" if torch.cuda.is_available() else "cpu").manual_seed(seed)
74
- else:
75
  seed = random.randint(1, 1000000000)
76
- generator = torch.Generator(device="cuda" if torch.cuda.is_available() else "cpu").manual_seed(seed)
77
 
78
- # Map sampler names to Diffusers scheduler names
79
- sampler_map = {
80
- "DPM++ 2M Karras": "dpmsolver++",
81
- "DPM++ SDE Karras": "dpmsolver++",
82
- "Euler": "euler",
83
- "Euler a": "euler_a",
84
- "Heun": "heun",
85
- "DDIM": "ddim"
86
  }
87
 
88
  try:
89
- # Generate image with LoRA
90
- image = pipe(
91
  prompt=prompt,
92
- negative_prompt=is_negative if is_negative else None,
93
  num_inference_steps=steps,
94
  guidance_scale=cfg_scale,
95
- generator=generator,
96
- strength=strength,
97
  width=width,
98
  height=height,
99
- cross_attention_kwargs={"scale": 0.8}, # LoRA strength adjustment
100
- ).images[0]
 
 
101
 
102
- png_img = convert_to_png(image)
103
- print(f'\033[1mGeneration {key} completed as PNG!\033[0m')
104
  return png_img
105
 
106
  except Exception as e:
107
- print(f"Generation error: {e}")
108
- raise gr.Error(f"Image generation failed: {str(e)}")
109
 
110
- # Light theme CSS (same as before)
111
  css = """
112
  #app-container {
113
  max-width: 800px;
@@ -139,7 +152,7 @@ h1 {
139
  """
140
 
141
  with gr.Blocks(theme=gr.themes.Default(primary_hue="green"), css=css) as app:
142
- gr.HTML("<center><h1>FLUX.1-Dev with LoRA (PNG Output)</h1></center>")
143
 
144
  with gr.Column(elem_id="app-container"):
145
  with gr.Row():
@@ -165,7 +178,6 @@ with gr.Blocks(theme=gr.themes.Default(primary_hue="green"), css=css) as app:
165
  steps = gr.Slider(35, label="Steps", minimum=10, maximum=100, step=1)
166
  cfg = gr.Slider(7.0, label="CFG Scale", minimum=1.0, maximum=20.0, step=0.5)
167
  with gr.Row():
168
- strength = gr.Slider(0.7, label="Strength", minimum=0.1, maximum=1.0, step=0.01)
169
  seed = gr.Number(-1, label="Seed (-1 for random)")
170
  method = gr.Radio(
171
  ["DPM++ 2M Karras", "DPM++ SDE Karras", "Euler", "Euler a", "Heun", "DDIM"],
@@ -184,8 +196,8 @@ with gr.Blocks(theme=gr.themes.Default(primary_hue="green"), css=css) as app:
184
  )
185
 
186
  generate_btn.click(
187
- fn=query,
188
- inputs=[text_prompt, negative_prompt, steps, cfg, method, seed, strength, width, height],
189
  outputs=output_image
190
  )
191
 
 
2
  import torch
3
  import random
4
  import os
5
+ import io
6
  from PIL import Image
7
  from deep_translator import GoogleTranslator
8
  from diffusers import DiffusionPipeline
9
+ from huggingface_hub import hf_hub_download, login
10
 
11
+ # Autentikasi Hugging Face
12
+ HF_TOKEN = os.getenv("HF_READ_TOKEN") # Ganti dengan token Anda atau set env variable
13
+ login(token=HF_TOKEN)
14
 
15
+ # Konfigurasi Model
16
  BASE_MODEL = "black-forest-labs/FLUX.1-dev"
17
  LORA_REPO = "burhansyam/uncen"
18
+ LORA_WEIGHTS_NAME = "uncen.safetensors" # Ganti jika nama file berbeda
19
+ torch_dtype = torch.float16 # Gunakan float16 untuk kompatibilitas lebih luas
20
 
21
+ # Inisialisasi Pipeline dengan LoRA
22
  def init_pipeline():
23
+ # Muat model dasar
24
  pipe = DiffusionPipeline.from_pretrained(
25
  BASE_MODEL,
26
+ torch_dtype=torch_dtype,
27
+ use_auth_token=HF_TOKEN
28
  )
29
 
30
+ # Muat weights LoRA
31
+ lora_path = hf_hub_download(
32
+ repo_id=LORA_REPO,
33
+ filename=LORA_WEIGHTS_NAME,
34
+ token=HF_TOKEN
35
  )
36
+ pipe.load_lora_weights(lora_path, adapter_name="uncen")
37
 
38
+ # Optimasi GPU jika tersedia
39
  if torch.cuda.is_available():
40
  pipe.to("cuda")
41
+ try:
42
+ pipe.enable_xformers_memory_efficient_attention()
43
+ except:
44
+ print("Xformers tidak tersedia, melanjutkan tanpa optimasi")
45
 
46
  return pipe
47
 
48
+ # Inisialisasi pipeline
49
+ try:
50
+ pipe = init_pipeline()
51
+ except Exception as e:
52
+ raise gr.Error(f"Gagal memuat model: {str(e)}. Pastikan token akses valid dan Anda memiliki izin.")
53
 
54
  def convert_to_png(image):
55
+ """Konversi gambar ke format PNG"""
56
  png_buffer = io.BytesIO()
57
  if image.mode == 'RGBA':
58
  image.save(png_buffer, format='PNG', optimize=True)
 
63
  png_buffer.seek(0)
64
  return Image.open(png_buffer)
65
 
66
+ def generate_image(
67
+ prompt,
68
+ negative_prompt="",
69
+ steps=35,
70
+ cfg_scale=7,
71
+ sampler="DPM++ 2M Karras",
72
+ seed=-1,
73
+ width=1024,
74
+ height=1024
75
+ ):
76
  if not prompt:
77
  return None
78
 
79
+ # Terjemahkan prompt jika bahasa Indonesia
 
 
80
  try:
81
+ translated_prompt = GoogleTranslator(source='id', target='en').translate(prompt)
82
+ prompt = f"{translated_prompt} | ultra detail, ultra quality, masterpiece"
83
+ except:
84
+ prompt = f"{prompt} | ultra detail, ultra quality, masterpiece"
 
85
 
86
+ # Set generator dengan seed
87
+ generator = torch.Generator(device="cuda" if torch.cuda.is_available() else "cpu")
88
+ if seed == -1:
 
 
 
 
89
  seed = random.randint(1, 1000000000)
90
+ generator.manual_seed(seed)
91
 
92
+ # Pemetaan sampler ke scheduler
93
+ sampler_config = {
94
+ "DPM++ 2M Karras": {"use_karras_sigmas": True},
95
+ "DPM++ SDE Karras": {"use_karras_sigmas": True},
96
+ "Euler": {},
97
+ "Euler a": {"use_karras_sigmas": False},
98
+ "Heun": {},
99
+ "DDIM": {}
100
  }
101
 
102
  try:
103
+ # Generate gambar dengan LoRA
104
+ result = pipe(
105
  prompt=prompt,
106
+ negative_prompt=negative_prompt,
107
  num_inference_steps=steps,
108
  guidance_scale=cfg_scale,
 
 
109
  width=width,
110
  height=height,
111
+ generator=generator,
112
+ cross_attention_kwargs={"scale": 1.0}, # Kontrol kekuatan LoRA
113
+ **sampler_config[sampler]
114
+ )
115
 
116
+ # Konversi ke PNG
117
+ png_img = convert_to_png(result.images[0])
118
  return png_img
119
 
120
  except Exception as e:
121
+ raise gr.Error(f"Generasi gambar gagal: {str(e)}")
 
122
 
123
+ # Antarmuka Gradio
124
  css = """
125
  #app-container {
126
  max-width: 800px;
 
152
  """
153
 
154
  with gr.Blocks(theme=gr.themes.Default(primary_hue="green"), css=css) as app:
155
+ gr.HTML("<center><h1>FLUX.1-Dev with Uncensored LoRA</h1></center>")
156
 
157
  with gr.Column(elem_id="app-container"):
158
  with gr.Row():
 
178
  steps = gr.Slider(35, label="Steps", minimum=10, maximum=100, step=1)
179
  cfg = gr.Slider(7.0, label="CFG Scale", minimum=1.0, maximum=20.0, step=0.5)
180
  with gr.Row():
 
181
  seed = gr.Number(-1, label="Seed (-1 for random)")
182
  method = gr.Radio(
183
  ["DPM++ 2M Karras", "DPM++ SDE Karras", "Euler", "Euler a", "Heun", "DDIM"],
 
196
  )
197
 
198
  generate_btn.click(
199
+ fn=generate_image,
200
+ inputs=[text_prompt, negative_prompt, steps, cfg, method, seed, width, height],
201
  outputs=output_image
202
  )
203