azaduddin commited on
Commit
376df9e
·
1 Parent(s): adf4aaa

copying the code from personal HF repo

Files changed (4) hide show
  1. README.md +4 -4
  2. app.py +152 -0
  3. live_preview_helpers.py +166 -0
  4. requirements.txt +14 -0
README.md CHANGED
@@ -1,10 +1,10 @@
1
  ---
2
  title: Flux Crypto Token
3
- emoji: 👁
4
- colorFrom: yellow
5
- colorTo: pink
6
  sdk: gradio
7
- sdk_version: 5.35.0
8
  app_file: app.py
9
  pinned: false
10
  ---
 
1
  ---
2
  title: Flux Crypto Token
3
+ emoji: 🦀
4
+ colorFrom: blue
5
+ colorTo: yellow
6
  sdk: gradio
7
+ sdk_version: 5.33.2
8
  app_file: app.py
9
  pinned: false
10
  ---
app.py ADDED
@@ -0,0 +1,152 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import numpy as np
3
+ import random
4
+ import spaces
5
+ import torch
6
+ from diffusers import DiffusionPipeline, FlowMatchEulerDiscreteScheduler, AutoencoderTiny, AutoencoderKL
7
+ from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast
8
+ from live_preview_helpers import calculate_shift, retrieve_timesteps, flux_pipe_call_that_returns_an_iterable_of_images
9
+ from huggingface_hub import login
10
+ import os
11
+ import tempfile
12
+ from peft import PeftModel
13
+
14
+ login(token=os.getenv('hf_token'))
15
+ dtype = torch.bfloat16
16
+ device = "cuda" if torch.cuda.is_available() else "cpu"
17
+
18
+ vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=dtype).to(device)
19
+ pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=dtype, vae=vae).to(device)
20
+ pipe.load_lora_weights("azad-uddin/meme-token", weight_name="meme-token.safetensors")
21
+ torch.cuda.empty_cache()
22
+
23
+ MAX_SEED = np.iinfo(np.int32).max
24
+ MAX_IMAGE_SIZE = 1024
25
+
26
+ pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)
27
+
28
+ @spaces.GPU(duration=75)
29
+ def infer(prompt, seed=42, randomize_seed=False, width=512, height=512, guidance_scale=3.5, num_inference_steps=28, progress=gr.Progress(track_tqdm=True)):
30
+ if randomize_seed:
31
+ seed = random.randint(0, MAX_SEED)
32
+ generator = torch.Generator().manual_seed(seed)
33
+
34
+ final_img = None
35
+ for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
36
+ prompt= f"memestylx, {prompt}, all within a circular emblem",
37
+ guidance_scale=guidance_scale,
38
+ num_inference_steps=num_inference_steps,
39
+ width=width,
40
+ height=height,
41
+ generator=generator,
42
+ output_type="pil",
43
+ good_vae=vae,
44
+ ):
45
+ final_img = img
46
+
47
+ if final_img is not None:
48
+ with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_file:
49
+ final_img.save(tmp_file.name, format="PNG")
50
+ image_path = tmp_file.name
51
+ return image_path
52
+ return None
53
+
54
+ examples = [
55
+ "flappy bird, takes you to height",
56
+ "spiderToken, make you fly",
57
+ "redbullCoin, gives you wiingss",
58
+ ]
59
+
60
+ css="""
61
+ #col-container {
62
+ margin: 0 auto;
63
+ max-width: 520px;
64
+ }
65
+ """
66
+
67
+ with gr.Blocks(css=css) as demo:
68
+
69
+ with gr.Column(elem_id="col-container"):
70
+ gr.Markdown(f"""# FLUX.1 [dev]
71
+ 12B param rectified flow transformer guidance-distilled from [FLUX.1 [pro]](https://blackforestlabs.ai/)
72
+ [[non-commercial license](https://huggingface.co/black-forest-labs/FLUX.1-dev/blob/main/LICENSE.md)] [[blog](https://blackforestlabs.ai/announcing-black-forest-labs/)] [[model](https://huggingface.co/black-forest-labs/FLUX.1-dev)]
73
+ """)
74
+
75
+ with gr.Row():
76
+
77
+ prompt = gr.Text(
78
+ label="Prompt",
79
+ show_label=False,
80
+ max_lines=1,
81
+ placeholder="Enter your prompt",
82
+ container=False,
83
+ )
84
+
85
+ run_button = gr.Button("Run", scale=0)
86
+
87
+ result = gr.Image(label="Result", show_label=False, type="filepath")
88
+
89
+ with gr.Accordion("Advanced Settings", open=False):
90
+
91
+ seed = gr.Slider(
92
+ label="Seed",
93
+ minimum=0,
94
+ maximum=MAX_SEED,
95
+ step=1,
96
+ value=0,
97
+ )
98
+
99
+ randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
100
+
101
+ with gr.Row():
102
+
103
+ width = gr.Slider(
104
+ label="Width",
105
+ minimum=256,
106
+ maximum=MAX_IMAGE_SIZE,
107
+ step=32,
108
+ value=1024,
109
+ )
110
+
111
+ height = gr.Slider(
112
+ label="Height",
113
+ minimum=256,
114
+ maximum=MAX_IMAGE_SIZE,
115
+ step=32,
116
+ value=1024,
117
+ )
118
+
119
+ with gr.Row():
120
+
121
+ guidance_scale = gr.Slider(
122
+ label="Guidance Scale",
123
+ minimum=1,
124
+ maximum=15,
125
+ step=0.1,
126
+ value=3.5,
127
+ )
128
+
129
+ num_inference_steps = gr.Slider(
130
+ label="Number of inference steps",
131
+ minimum=1,
132
+ maximum=50,
133
+ step=1,
134
+ value=28,
135
+ )
136
+
137
+ gr.Examples(
138
+ examples = examples,
139
+ fn = infer,
140
+ inputs = [prompt],
141
+ outputs = [result],
142
+ cache_examples="lazy"
143
+ )
144
+
145
+ gr.on(
146
+ triggers=[run_button.click, prompt.submit],
147
+ fn = infer,
148
+ inputs = [prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
149
+ outputs = [result]
150
+ )
151
+
152
+ demo.launch()
live_preview_helpers.py ADDED
@@ -0,0 +1,166 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import numpy as np
3
+ from diffusers import FluxPipeline, AutoencoderTiny, FlowMatchEulerDiscreteScheduler
4
+ from typing import Any, Dict, List, Optional, Union
5
+
6
+ # Helper functions
7
+ def calculate_shift(
8
+ image_seq_len,
9
+ base_seq_len: int = 256,
10
+ max_seq_len: int = 4096,
11
+ base_shift: float = 0.5,
12
+ max_shift: float = 1.16,
13
+ ):
14
+ m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
15
+ b = base_shift - m * base_seq_len
16
+ mu = image_seq_len * m + b
17
+ return mu
18
+
19
+ def retrieve_timesteps(
20
+ scheduler,
21
+ num_inference_steps: Optional[int] = None,
22
+ device: Optional[Union[str, torch.device]] = None,
23
+ timesteps: Optional[List[int]] = None,
24
+ sigmas: Optional[List[float]] = None,
25
+ **kwargs,
26
+ ):
27
+ if timesteps is not None and sigmas is not None:
28
+ raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
29
+ if timesteps is not None:
30
+ scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
31
+ timesteps = scheduler.timesteps
32
+ num_inference_steps = len(timesteps)
33
+ elif sigmas is not None:
34
+ scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
35
+ timesteps = scheduler.timesteps
36
+ num_inference_steps = len(timesteps)
37
+ else:
38
+ scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
39
+ timesteps = scheduler.timesteps
40
+ return timesteps, num_inference_steps
41
+
42
+ # FLUX pipeline function
43
+ @torch.inference_mode()
44
+ def flux_pipe_call_that_returns_an_iterable_of_images(
45
+ self,
46
+ prompt: Union[str, List[str]] = None,
47
+ prompt_2: Optional[Union[str, List[str]]] = None,
48
+ height: Optional[int] = None,
49
+ width: Optional[int] = None,
50
+ num_inference_steps: int = 28,
51
+ timesteps: List[int] = None,
52
+ guidance_scale: float = 3.5,
53
+ num_images_per_prompt: Optional[int] = 1,
54
+ generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
55
+ latents: Optional[torch.FloatTensor] = None,
56
+ prompt_embeds: Optional[torch.FloatTensor] = None,
57
+ pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
58
+ output_type: Optional[str] = "pil",
59
+ return_dict: bool = True,
60
+ joint_attention_kwargs: Optional[Dict[str, Any]] = None,
61
+ max_sequence_length: int = 512,
62
+ good_vae: Optional[Any] = None,
63
+ ):
64
+ height = height or self.default_sample_size * self.vae_scale_factor
65
+ width = width or self.default_sample_size * self.vae_scale_factor
66
+
67
+ # 1. Check inputs
68
+ self.check_inputs(
69
+ prompt,
70
+ prompt_2,
71
+ height,
72
+ width,
73
+ prompt_embeds=prompt_embeds,
74
+ pooled_prompt_embeds=pooled_prompt_embeds,
75
+ max_sequence_length=max_sequence_length,
76
+ )
77
+
78
+ self._guidance_scale = guidance_scale
79
+ self._joint_attention_kwargs = joint_attention_kwargs
80
+ self._interrupt = False
81
+
82
+ # 2. Define call parameters
83
+ batch_size = 1 if isinstance(prompt, str) else len(prompt)
84
+ device = self._execution_device
85
+
86
+ # 3. Encode prompt
87
+ lora_scale = joint_attention_kwargs.get("scale", None) if joint_attention_kwargs is not None else None
88
+ prompt_embeds, pooled_prompt_embeds, text_ids = self.encode_prompt(
89
+ prompt=prompt,
90
+ prompt_2=prompt_2,
91
+ prompt_embeds=prompt_embeds,
92
+ pooled_prompt_embeds=pooled_prompt_embeds,
93
+ device=device,
94
+ num_images_per_prompt=num_images_per_prompt,
95
+ max_sequence_length=max_sequence_length,
96
+ lora_scale=lora_scale,
97
+ )
98
+ # 4. Prepare latent variables
99
+ num_channels_latents = self.transformer.config.in_channels // 4
100
+ latents, latent_image_ids = self.prepare_latents(
101
+ batch_size * num_images_per_prompt,
102
+ num_channels_latents,
103
+ height,
104
+ width,
105
+ prompt_embeds.dtype,
106
+ device,
107
+ generator,
108
+ latents,
109
+ )
110
+ # 5. Prepare timesteps
111
+ sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)
112
+ image_seq_len = latents.shape[1]
113
+ mu = calculate_shift(
114
+ image_seq_len,
115
+ self.scheduler.config.base_image_seq_len,
116
+ self.scheduler.config.max_image_seq_len,
117
+ self.scheduler.config.base_shift,
118
+ self.scheduler.config.max_shift,
119
+ )
120
+ timesteps, num_inference_steps = retrieve_timesteps(
121
+ self.scheduler,
122
+ num_inference_steps,
123
+ device,
124
+ timesteps,
125
+ sigmas,
126
+ mu=mu,
127
+ )
128
+ self._num_timesteps = len(timesteps)
129
+
130
+ # Handle guidance
131
+ guidance = torch.full([1], guidance_scale, device=device, dtype=torch.float32).expand(latents.shape[0]) if self.transformer.config.guidance_embeds else None
132
+
133
+ # 6. Denoising loop
134
+ for i, t in enumerate(timesteps):
135
+ if self.interrupt:
136
+ continue
137
+
138
+ timestep = t.expand(latents.shape[0]).to(latents.dtype)
139
+
140
+ noise_pred = self.transformer(
141
+ hidden_states=latents,
142
+ timestep=timestep / 1000,
143
+ guidance=guidance,
144
+ pooled_projections=pooled_prompt_embeds,
145
+ encoder_hidden_states=prompt_embeds,
146
+ txt_ids=text_ids,
147
+ img_ids=latent_image_ids,
148
+ joint_attention_kwargs=self.joint_attention_kwargs,
149
+ return_dict=False,
150
+ )[0]
151
+ # Yield intermediate result
152
+ latents_for_image = self._unpack_latents(latents, height, width, self.vae_scale_factor)
153
+ latents_for_image = (latents_for_image / self.vae.config.scaling_factor) + self.vae.config.shift_factor
154
+ image = self.vae.decode(latents_for_image, return_dict=False)[0]
155
+ yield self.image_processor.postprocess(image, output_type=output_type)[0]
156
+
157
+ latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
158
+ torch.cuda.empty_cache()
159
+
160
+ # Final image using good_vae
161
+ latents = self._unpack_latents(latents, height, width, self.vae_scale_factor)
162
+ latents = (latents / good_vae.config.scaling_factor) + good_vae.config.shift_factor
163
+ image = good_vae.decode(latents, return_dict=False)[0]
164
+ self.maybe_free_model_hooks()
165
+ torch.cuda.empty_cache()
166
+ yield self.image_processor.postprocess(image, output_type=output_type)[0]
requirements.txt ADDED
@@ -0,0 +1,14 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ git+https://github.com/huggingface/diffusers.git
2
+ git+https://github.com/huggingface/peft.git
3
+ gradio==4.29.0 # Explicitly pin Gradio to a known stable version
4
+ accelerate
5
+ diffusers
6
+ torch
7
+ numpy
8
+ transformers
9
+ xformers
10
+ sentencepiece
11
+ scipy
12
+ pydantic==2.10.6 # Pinning pydantic to avoid potential conflicts with Gradio
13
+ Pillow
14
+ spaces