prithivMLmods commited on
Commit
5bdb5dc
·
verified ·
1 Parent(s): 95af69c

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +23 -12
app.py CHANGED
@@ -5,21 +5,24 @@ from PIL import Image
5
  from diffusers import DiffusionPipeline
6
  import random
7
  import uuid
 
8
  import numpy as np
9
  import time
10
  import zipfile
11
  import os
12
 
13
  # Description for the app
14
- DESCRIPTION = """## Qwen Image HPC/."""
15
 
16
  # Helper functions
17
  def save_image(img):
 
18
  unique_name = str(uuid.uuid4()) + ".png"
19
  img.save(unique_name)
20
  return unique_name
21
 
22
  def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
 
23
  if randomize_seed:
24
  seed = random.randint(0, MAX_SEED)
25
  return seed
@@ -27,24 +30,32 @@ def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
27
  MAX_SEED = np.iinfo(np.int32).max
28
  MAX_IMAGE_SIZE = 2048
29
 
30
- # Load Qwen/Qwen-Image pipeline with regional compilation
31
  dtype = torch.bfloat16
32
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
33
 
34
- # --- Model Loading with Regional Compilation ---
35
- ckpt_id = "Qwen/Qwen-Image"
36
  pipe_qwen = DiffusionPipeline.from_pretrained(
37
- ckpt_id, torch_dtype=dtype
 
38
  ).to(device)
39
- pipe_qwen.transformer.compile_repeated_blocks(fullgraph=True)
 
 
 
 
 
 
 
40
 
41
  # Aspect ratios
42
  aspect_ratios = {
43
- "1:1": (1328, 1328),
44
- "16:9": (1664, 928),
45
- "9:16": (928, 1664),
46
- "4:3": (1472, 1140),
47
- "3:4": (1140, 1472)
48
  }
49
 
50
  # Generation function for Qwen/Qwen-Image
@@ -281,4 +292,4 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
281
  )
282
 
283
  if __name__ == "__main__":
284
- demo.queue(max_size=50).launch(share=False, mcp_server=True, ssr_mode=False, show_error=True)
 
5
  from diffusers import DiffusionPipeline
6
  import random
7
  import uuid
8
+ from typing import Union, List, Optional
9
  import numpy as np
10
  import time
11
  import zipfile
12
  import os
13
 
14
  # Description for the app
15
+ DESCRIPTION = """## Qwen Image Hpc/."""
16
 
17
  # Helper functions
18
  def save_image(img):
19
+ """Saves a PIL image to a file with a unique name."""
20
  unique_name = str(uuid.uuid4()) + ".png"
21
  img.save(unique_name)
22
  return unique_name
23
 
24
  def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
25
+ """Generates a random seed if the randomize option is enabled."""
26
  if randomize_seed:
27
  seed = random.randint(0, MAX_SEED)
28
  return seed
 
30
  MAX_SEED = np.iinfo(np.int32).max
31
  MAX_IMAGE_SIZE = 2048
32
 
33
+ # Load Qwen/Qwen-Image pipeline
34
  dtype = torch.bfloat16
35
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
36
 
37
+ # --- Model Loading ---
38
+ # Load the pipeline from pretrained weights
39
  pipe_qwen = DiffusionPipeline.from_pretrained(
40
+ "Qwen/Qwen-Image",
41
+ torch_dtype=dtype
42
  ).to(device)
43
+
44
+ # --- Regional Compilation ---
45
+ # Apply regional compilation to speed up cold-starts while retaining benefits.
46
+ # This compiles the repeated transformer blocks for ~2x faster initialization.
47
+ print("Applying regional compilation...")
48
+ pipe_qwen.transformer.compile(fullgraph=True, mode="reduce-overhead")
49
+ print("Compilation complete.")
50
+
51
 
52
  # Aspect ratios
53
  aspect_ratios = {
54
+ "1:1": (1024, 1024),
55
+ "16:9": (1344, 768),
56
+ "9:16": (768, 1344),
57
+ "4:3": (1152, 896),
58
+ "3:4": (896, 1152)
59
  }
60
 
61
  # Generation function for Qwen/Qwen-Image
 
292
  )
293
 
294
  if __name__ == "__main__":
295
+ demo.queue(max_size=50).launch(share=False)