blanchon commited on
Commit
d34889d
·
1 Parent(s): 0628786
Files changed (1) hide show
  1. app-full.py +12 -10
app-full.py CHANGED
@@ -19,13 +19,13 @@ MODEL_CONFIGS = {
19
 
20
  # Supported image sizes
21
  RESOLUTION_OPTIONS: list[str] = [
22
- "1024 x 1024 (Square)",
23
- "768 x 1360 (Portrait)",
24
- "1360 x 768 (Landscape)",
25
- "880 x 1168 (Portrait)",
26
- "1168 x 880 (Landscape)",
27
- "1248 x 832 (Landscape)",
28
- "832 x 1248 (Portrait)",
29
  ]
30
 
31
 
@@ -60,13 +60,15 @@ pipe = HiDreamImagePipeline.from_pretrained(
60
  pipe.transformer = transformer
61
 
62
 
63
- @spaces.GPU(duration=60 * 3)
64
  def generate_image(
65
  prompt: str, resolution: str, seed: int, progress=gr.Progress(track_tqdm=True)
66
  ) -> tuple[PIL.Image.Image, int]:
67
  if seed == -1:
68
  seed = torch.randint(0, 1_000_000, (1,)).item()
69
 
 
 
70
  height, width = tuple(map(int, resolution.replace(" ", "").split("x")))
71
  generator = torch.Generator("cuda").manual_seed(seed)
72
 
@@ -84,8 +86,8 @@ def generate_image(
84
 
85
 
86
  # Gradio UI
87
- with gr.Blocks(title="HiDream Image Generator") as demo:
88
- gr.Markdown("## 🌈 HiDream Image Generator")
89
 
90
  with gr.Row():
91
  with gr.Column():
 
19
 
20
  # Supported image sizes
21
  RESOLUTION_OPTIONS: list[str] = [
22
+ "1024 x 1024",
23
+ "768 x 1360",
24
+ "1360 x 768",
25
+ "880 x 1168",
26
+ "1168 x 880",
27
+ "1248 x 832",
28
+ "832 x 1248",
29
  ]
30
 
31
 
 
60
  pipe.transformer = transformer
61
 
62
 
63
+ @spaces.GPU(duration=120)
64
  def generate_image(
65
  prompt: str, resolution: str, seed: int, progress=gr.Progress(track_tqdm=True)
66
  ) -> tuple[PIL.Image.Image, int]:
67
  if seed == -1:
68
  seed = torch.randint(0, 1_000_000, (1,)).item()
69
 
70
+ msg = "ℹ️ This spaces currently crash because of the memory usage. Please help me fix 😅"
71
+ raise gr.Error(msg, duration=10)
72
  height, width = tuple(map(int, resolution.replace(" ", "").split("x")))
73
  generator = torch.Generator("cuda").manual_seed(seed)
74
 
 
86
 
87
 
88
  # Gradio UI
89
+ with gr.Blocks(title="HiDream Image Generator Full") as demo:
90
+ gr.Markdown("## 🌈 HiDream Image Generator Full")
91
 
92
  with gr.Row():
93
  with gr.Column():