ovi054 commited on
Commit
22ada24
·
verified ·
1 Parent(s): 2875916

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +15 -23
app.py CHANGED
@@ -8,7 +8,7 @@ import numpy as np
8
 
9
  # Set data type
10
  dtype = torch.bfloat16
11
- device = "cpu" # Initial device for model loading; inference will use GPU
12
 
13
  # Load tokenizer and text encoder for Llama
14
  try:
@@ -25,12 +25,12 @@ except Exception as e:
25
  # Load the HiDreamImagePipeline
26
  try:
27
  pipe = HiDreamImagePipeline.from_pretrained(
28
- "HiDream-ai/HiDream-I1-Full",
29
  tokenizer_4=tokenizer_4,
30
  text_encoder_4=text_encoder_4,
31
  torch_dtype=dtype,
32
  ).to(device)
33
- pipe.enable_model_cpu_offload() # Offload to CPU when not in use, critical for Spaces
34
  except Exception as e:
35
  raise Exception(f"Failed to load HiDreamImagePipeline: {e}. Ensure you have access to 'HiDream-ai/HiDream-I1-Full'.")
36
 
@@ -40,16 +40,13 @@ MAX_IMAGE_SIZE = 2048
40
 
41
  # Inference function with GPU access
42
  @spaces.GPU()
43
- def infer(prompt, negative_prompt="", seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=50, guidance_scale=5.0, progress=gr.Progress(track_tqdm=True)):
44
- # Ensure the model is on GPU for inference
45
- pipe.to("cuda")
46
-
47
  try:
48
  if randomize_seed:
49
  seed = random.randint(0, MAX_SEED)
50
  generator = torch.Generator("cuda").manual_seed(seed)
51
 
52
- # Generate the image
53
  image = pipe(
54
  prompt=prompt,
55
  negative_prompt=negative_prompt,
@@ -61,13 +58,9 @@ def infer(prompt, negative_prompt="", seed=42, randomize_seed=False, width=1024,
61
  output_type="pil",
62
  ).images[0]
63
 
64
- # Clear GPU memory
65
- torch.cuda.empty_cache()
66
-
67
  return image, seed
68
  finally:
69
- # Move model back to CPU to free GPU memory
70
- pipe.to("cpu")
71
  torch.cuda.empty_cache()
72
 
73
  # Define examples
@@ -89,7 +82,7 @@ css = """
89
  color: white !important;
90
  }
91
  .generate-btn:hover {
92
- transform: translateY(-2px);
93
  box-shadow: 0 5px 15px rgba(0,0,0,0.2);
94
  }
95
  """
@@ -107,6 +100,12 @@ with gr.Blocks(css=css) as app:
107
  lines=3,
108
  elem_id="prompt-text-input"
109
  )
 
 
 
 
 
 
110
  with gr.Row():
111
  with gr.Accordion("Advanced Settings", open=False):
112
  with gr.Row():
@@ -127,14 +126,14 @@ with gr.Blocks(css=css) as app:
127
  with gr.Row():
128
  steps = gr.Slider(
129
  label="Inference Steps",
130
- value=50,
131
  minimum=1,
132
  maximum=100,
133
  step=1
134
  )
135
  cfg = gr.Slider(
136
  label="Guidance Scale",
137
- value=5.0,
138
  minimum=1,
139
  maximum=20,
140
  step=0.5
@@ -151,12 +150,6 @@ with gr.Blocks(css=css) as app:
151
  label="Randomize Seed",
152
  value=True
153
  )
154
- with gr.Row():
155
- negative_prompt = gr.Textbox(
156
- label="Negative Prompt",
157
- placeholder="Enter what to avoid (optional)",
158
- lines=2
159
- )
160
  with gr.Row():
161
  text_button = gr.Button(
162
  "✨ Generate Image",
@@ -189,5 +182,4 @@ with gr.Blocks(css=css) as app:
189
  outputs=[image_output, seed_output]
190
  )
191
 
192
- # Launch the app
193
  app.launch(share=True)
 
8
 
9
  # Set data type
10
  dtype = torch.bfloat16
11
+ device = "cpu" # Use CPU for model loading to avoid CUDA initialization
12
 
13
  # Load tokenizer and text encoder for Llama
14
  try:
 
25
  # Load the HiDreamImagePipeline
26
  try:
27
  pipe = HiDreamImagePipeline.from_pretrained(
28
+ "HiDream-ai/HiDream-I1-Dev",
29
  tokenizer_4=tokenizer_4,
30
  text_encoder_4=text_encoder_4,
31
  torch_dtype=dtype,
32
  ).to(device)
33
+ pipe.enable_model_cpu_offload() # Offload to CPU, automatically manages GPU placement
34
  except Exception as e:
35
  raise Exception(f"Failed to load HiDreamImagePipeline: {e}. Ensure you have access to 'HiDream-ai/HiDream-I1-Full'.")
36
 
 
40
 
41
  # Inference function with GPU access
42
  @spaces.GPU()
43
+ def infer(prompt, negative_prompt="", seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=28, guidance_scale=3.5, progress=gr.Progress(track_tqdm=True)):
 
 
 
44
  try:
45
  if randomize_seed:
46
  seed = random.randint(0, MAX_SEED)
47
  generator = torch.Generator("cuda").manual_seed(seed)
48
 
49
+ # Generate the image; offloading handles device placement
50
  image = pipe(
51
  prompt=prompt,
52
  negative_prompt=negative_prompt,
 
58
  output_type="pil",
59
  ).images[0]
60
 
 
 
 
61
  return image, seed
62
  finally:
63
+ # Clear GPU memory
 
64
  torch.cuda.empty_cache()
65
 
66
  # Define examples
 
82
  color: white !important;
83
  }
84
  .generate-btn:hover {
85
+ transform: translateY2px);
86
  box-shadow: 0 5px 15px rgba(0,0,0,0.2);
87
  }
88
  """
 
100
  lines=3,
101
  elem_id="prompt-text-input"
102
  )
103
+ with gr.Row():
104
+ negative_prompt = gr.Textbox(
105
+ label="Negative Prompt",
106
+ placeholder="Enter what to avoid (optional)",
107
+ lines=2
108
+ )
109
  with gr.Row():
110
  with gr.Accordion("Advanced Settings", open=False):
111
  with gr.Row():
 
126
  with gr.Row():
127
  steps = gr.Slider(
128
  label="Inference Steps",
129
+ value=28,
130
  minimum=1,
131
  maximum=100,
132
  step=1
133
  )
134
  cfg = gr.Slider(
135
  label="Guidance Scale",
136
+ value=3.5,
137
  minimum=1,
138
  maximum=20,
139
  step=0.5
 
150
  label="Randomize Seed",
151
  value=True
152
  )
 
 
 
 
 
 
153
  with gr.Row():
154
  text_button = gr.Button(
155
  "✨ Generate Image",
 
182
  outputs=[image_output, seed_output]
183
  )
184
 
 
185
  app.launch(share=True)