Spaces:
Running
Running
Upload 7 files
Browse files- app.py +273 -132
- download_models.py +79 -15
- error_page.html +99 -0
- nltk_setup.py +16 -0
- requirements.txt +3 -9
- startup.sh +12 -3
app.py
CHANGED
@@ -10,43 +10,57 @@ import gradio as gr
|
|
10 |
from omegaconf import OmegaConf
|
11 |
from scipy.stats import truncnorm
|
12 |
import subprocess
|
|
|
|
|
|
|
|
|
|
|
13 |
|
14 |
# First run the download_models.py script if models haven't been downloaded
|
15 |
-
if not os.path.exists('data/state_epoch_1220.pth') or not os.path.exists('data/text_encoder200.pth'):
|
16 |
print("Downloading necessary model files...")
|
17 |
try:
|
18 |
subprocess.check_call([sys.executable, "download_models.py"])
|
19 |
except subprocess.CalledProcessError as e:
|
20 |
print(f"Error downloading models: {e}")
|
21 |
-
print("Please
|
22 |
|
23 |
-
#
|
24 |
-
|
|
|
|
|
25 |
|
26 |
-
# Import necessary modules from the DF-GAN code
|
27 |
-
from models.DAMSM import RNN_ENCODER
|
28 |
-
from models.GAN import NetG
|
|
|
|
|
|
|
29 |
|
30 |
# Utility functions
|
31 |
def load_model_weights(model, weights, multi_gpus=False, train=False):
|
32 |
"""Load model weights with proper handling of module prefix"""
|
33 |
-
|
34 |
-
|
35 |
-
|
36 |
-
|
37 |
-
|
38 |
-
|
39 |
-
if
|
40 |
-
|
41 |
-
|
42 |
-
|
43 |
-
|
|
|
|
|
|
|
44 |
else:
|
45 |
state_dict = weights
|
46 |
-
|
47 |
-
state_dict
|
48 |
-
|
49 |
-
|
|
|
50 |
return model
|
51 |
|
52 |
def get_tokenizer():
|
@@ -86,22 +100,32 @@ def tokenize_and_build_captions(input_text, wordtoix):
|
|
86 |
|
87 |
def encode_caption(caption, caption_len, text_encoder, device):
|
88 |
"""Encode caption using text encoder"""
|
89 |
-
|
90 |
-
|
91 |
-
|
92 |
-
|
93 |
-
|
94 |
-
|
|
|
|
|
|
|
|
|
|
|
95 |
|
96 |
def save_img(img_tensor):
|
97 |
"""Convert image tensor to PIL Image"""
|
98 |
-
|
99 |
-
|
100 |
-
|
101 |
-
|
102 |
-
|
103 |
-
|
104 |
-
|
|
|
|
|
|
|
|
|
|
|
105 |
|
106 |
# Load configuration
|
107 |
config = {
|
@@ -114,124 +138,241 @@ config = {
|
|
114 |
'trunc_rate': 0.88,
|
115 |
}
|
116 |
|
|
|
117 |
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
118 |
print(f"Using device: {device}")
|
119 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
120 |
# Load vocab and models
|
121 |
def load_models():
|
122 |
-
|
123 |
-
with open('data/captions_DAMSM.pickle', 'rb') as f:
|
124 |
-
x = pickle.load(f)
|
125 |
-
wordtoix = x[3]
|
126 |
-
ixtoword = x[2]
|
127 |
-
del x
|
128 |
-
|
129 |
-
# Initialize text encoder
|
130 |
-
text_encoder = RNN_ENCODER(len(wordtoix), nhidden=config['cond_dim'])
|
131 |
-
text_encoder_path = 'data/text_encoder200.pth'
|
132 |
-
state_dict = torch.load(text_encoder_path, map_location='cpu')
|
133 |
-
text_encoder = load_model_weights(text_encoder, state_dict)
|
134 |
-
text_encoder.to(device)
|
135 |
-
for p in text_encoder.parameters():
|
136 |
-
p.requires_grad = False
|
137 |
-
text_encoder.eval()
|
138 |
|
139 |
-
|
140 |
-
|
141 |
-
|
142 |
-
|
143 |
-
|
144 |
-
|
145 |
-
|
146 |
-
|
147 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
148 |
|
149 |
-
|
|
|
|
|
|
|
|
|
|
|
150 |
|
151 |
def generate_image(text_input, num_images=1, seed=None):
|
152 |
"""Generate images from text description"""
|
153 |
if not text_input.strip():
|
154 |
-
return [
|
155 |
-
|
156 |
-
cap_array, cap_len = tokenize_and_build_captions(text_input, wordtoix)
|
157 |
-
|
158 |
-
if cap_len == 0:
|
159 |
-
return [Image.new('RGB', (256, 256), color='red')] * num_images
|
160 |
|
161 |
-
|
162 |
-
|
163 |
-
|
164 |
-
|
165 |
-
|
166 |
-
|
167 |
-
|
168 |
-
|
169 |
-
|
170 |
-
|
171 |
-
|
172 |
-
|
173 |
-
|
174 |
-
|
175 |
-
|
176 |
-
|
177 |
-
|
178 |
-
|
179 |
-
|
180 |
-
|
181 |
-
|
182 |
-
|
183 |
-
|
184 |
-
|
185 |
-
|
186 |
-
|
187 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
188 |
|
189 |
# Create Gradio interface
|
190 |
def generate_images_interface(text, num_images, random_seed):
|
191 |
-
seed = int(random_seed) if random_seed else None
|
192 |
return generate_image(text, num_images, seed)
|
193 |
|
|
|
194 |
with gr.Blocks(title="Bird Image Generator") as demo:
|
195 |
-
|
196 |
-
|
197 |
-
|
198 |
-
|
199 |
-
|
200 |
-
|
201 |
-
|
202 |
-
|
203 |
-
|
204 |
-
|
205 |
-
|
206 |
-
|
207 |
-
|
208 |
-
|
209 |
-
|
210 |
-
|
211 |
-
|
212 |
-
|
213 |
-
|
214 |
-
|
215 |
-
|
216 |
-
|
217 |
-
|
218 |
-
|
219 |
-
|
220 |
-
|
221 |
-
|
222 |
-
"
|
223 |
-
|
224 |
-
|
225 |
-
|
226 |
-
|
227 |
-
|
228 |
-
|
229 |
-
|
230 |
-
|
231 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
232 |
|
233 |
# Launch the app with appropriate configurations for Hugging Face Spaces
|
234 |
if __name__ == "__main__":
|
|
|
|
|
|
|
235 |
demo.launch(
|
236 |
server_name="0.0.0.0", # Bind to all network interfaces
|
237 |
share=False, # Don't use share links
|
|
|
10 |
from omegaconf import OmegaConf
|
11 |
from scipy.stats import truncnorm
|
12 |
import subprocess
|
13 |
+
import traceback
|
14 |
+
import time
|
15 |
+
|
16 |
+
# Create a flag to track model loading status
|
17 |
+
models_loaded_successfully = False
|
18 |
|
19 |
# First run the download_models.py script if models haven't been downloaded
|
20 |
+
if not os.path.exists('data/state_epoch_1220.pth') or not os.path.exists('data/text_encoder200.pth') or not os.path.exists('data/captions_DAMSM.pickle'):
|
21 |
print("Downloading necessary model files...")
|
22 |
try:
|
23 |
subprocess.check_call([sys.executable, "download_models.py"])
|
24 |
except subprocess.CalledProcessError as e:
|
25 |
print(f"Error downloading models: {e}")
|
26 |
+
print("Please check the error message above. The application will attempt to continue with fallback settings.")
|
27 |
|
28 |
+
# Setup system paths
|
29 |
+
try:
|
30 |
+
# Add the code directory to the Python path
|
31 |
+
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "DF-GAN/code"))
|
32 |
|
33 |
+
# Import necessary modules from the DF-GAN code
|
34 |
+
from models.DAMSM import RNN_ENCODER
|
35 |
+
from models.GAN import NetG
|
36 |
+
except ImportError as e:
|
37 |
+
print(f"Error importing required modules: {e}")
|
38 |
+
print("The application may not function correctly.")
|
39 |
|
40 |
# Utility functions
|
41 |
def load_model_weights(model, weights, multi_gpus=False, train=False):
|
42 |
"""Load model weights with proper handling of module prefix"""
|
43 |
+
try:
|
44 |
+
if list(weights.keys())[0].find('module')==-1:
|
45 |
+
pretrained_with_multi_gpu = False
|
46 |
+
else:
|
47 |
+
pretrained_with_multi_gpu = True
|
48 |
+
|
49 |
+
if (multi_gpus==False) or (train==False):
|
50 |
+
if pretrained_with_multi_gpu:
|
51 |
+
state_dict = {
|
52 |
+
key[7:]: value
|
53 |
+
for key, value in weights.items()
|
54 |
+
}
|
55 |
+
else:
|
56 |
+
state_dict = weights
|
57 |
else:
|
58 |
state_dict = weights
|
59 |
+
|
60 |
+
model.load_state_dict(state_dict)
|
61 |
+
except Exception as e:
|
62 |
+
print(f"Error loading model weights: {e}")
|
63 |
+
print("Using model with random weights instead.")
|
64 |
return model
|
65 |
|
66 |
def get_tokenizer():
|
|
|
100 |
|
101 |
def encode_caption(caption, caption_len, text_encoder, device):
|
102 |
"""Encode caption using text encoder"""
|
103 |
+
try:
|
104 |
+
with torch.no_grad():
|
105 |
+
caption = torch.tensor([caption]).to(device)
|
106 |
+
caption_len = torch.tensor([caption_len]).to(device)
|
107 |
+
hidden = text_encoder.init_hidden(1)
|
108 |
+
_, sent_emb = text_encoder(caption, caption_len, hidden)
|
109 |
+
return sent_emb
|
110 |
+
except Exception as e:
|
111 |
+
print(f"Error encoding caption: {e}")
|
112 |
+
# Return a random embedding as fallback
|
113 |
+
return torch.randn(1, 256).to(device)
|
114 |
|
115 |
def save_img(img_tensor):
|
116 |
"""Convert image tensor to PIL Image"""
|
117 |
+
try:
|
118 |
+
im = img_tensor.data.cpu().numpy()
|
119 |
+
# [-1, 1] --> [0, 255]
|
120 |
+
im = (im + 1.0) * 127.5
|
121 |
+
im = im.astype(np.uint8)
|
122 |
+
im = np.transpose(im, (1, 2, 0))
|
123 |
+
im = Image.fromarray(im)
|
124 |
+
return im
|
125 |
+
except Exception as e:
|
126 |
+
print(f"Error converting image tensor to PIL Image: {e}")
|
127 |
+
# Return a red placeholder image as fallback
|
128 |
+
return Image.new('RGB', (256, 256), color='red')
|
129 |
|
130 |
# Load configuration
|
131 |
config = {
|
|
|
138 |
'trunc_rate': 0.88,
|
139 |
}
|
140 |
|
141 |
+
# Determine device
|
142 |
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
143 |
print(f"Using device: {device}")
|
144 |
|
145 |
+
# Global variables for models
|
146 |
+
wordtoix = {}
|
147 |
+
ixtoword = {}
|
148 |
+
text_encoder = None
|
149 |
+
netG = None
|
150 |
+
models_loaded = False
|
151 |
+
|
152 |
# Load vocab and models
|
153 |
def load_models():
|
154 |
+
global wordtoix, ixtoword, text_encoder, netG, models_loaded, models_loaded_successfully
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
155 |
|
156 |
+
try:
|
157 |
+
# Load vocabulary
|
158 |
+
if os.path.exists('data/captions_DAMSM.pickle'):
|
159 |
+
with open('data/captions_DAMSM.pickle', 'rb') as f:
|
160 |
+
x = pickle.load(f)
|
161 |
+
wordtoix = x[3]
|
162 |
+
ixtoword = x[2]
|
163 |
+
del x
|
164 |
+
else:
|
165 |
+
print("Warning: captions_DAMSM.pickle not found. Using fallback vocabulary.")
|
166 |
+
# Fallback vocabulary
|
167 |
+
wordtoix = {"the": 1, "bird": 2, "is": 3, "a": 4, "with": 5, "and": 6, "red": 7, "black": 8, "yellow": 9}
|
168 |
+
ixtoword = {v: k for k, v in wordtoix.items()}
|
169 |
+
|
170 |
+
# Initialize text encoder
|
171 |
+
text_encoder = RNN_ENCODER(len(wordtoix), nhidden=config['cond_dim'])
|
172 |
+
text_encoder_path = 'data/text_encoder200.pth'
|
173 |
+
if os.path.exists(text_encoder_path):
|
174 |
+
state_dict = torch.load(text_encoder_path, map_location='cpu')
|
175 |
+
text_encoder = load_model_weights(text_encoder, state_dict)
|
176 |
+
else:
|
177 |
+
print("Warning: text_encoder200.pth not found. Using random weights.")
|
178 |
+
text_encoder.to(device)
|
179 |
+
for p in text_encoder.parameters():
|
180 |
+
p.requires_grad = False
|
181 |
+
text_encoder.eval()
|
182 |
+
|
183 |
+
# Initialize generator
|
184 |
+
netG = NetG(config['nf'], config['z_dim'], config['cond_dim'], config['imsize'], config['ch_size'])
|
185 |
+
netG_path = 'data/state_epoch_1220.pth'
|
186 |
+
if os.path.exists(netG_path):
|
187 |
+
state_dict = torch.load(netG_path, map_location='cpu')
|
188 |
+
if 'model' in state_dict and 'netG' in state_dict['model']:
|
189 |
+
netG = load_model_weights(netG, state_dict['model']['netG'])
|
190 |
+
models_loaded_successfully = True
|
191 |
+
else:
|
192 |
+
print("Warning: state_epoch_1220.pth has unexpected format. Using random weights.")
|
193 |
+
else:
|
194 |
+
print("Warning: state_epoch_1220.pth not found. Using random weights.")
|
195 |
+
netG.to(device)
|
196 |
+
netG.eval()
|
197 |
+
|
198 |
+
models_loaded = True
|
199 |
+
return wordtoix, ixtoword, text_encoder, netG
|
200 |
+
except Exception as e:
|
201 |
+
print(f"Error loading models: {e}")
|
202 |
+
traceback.print_exc()
|
203 |
+
print("Using fallback models instead.")
|
204 |
+
|
205 |
+
# Fallback vocabulary
|
206 |
+
wordtoix = {"the": 1, "bird": 2, "is": 3, "a": 4, "with": 5, "and": 6, "red": 7, "black": 8, "yellow": 9}
|
207 |
+
ixtoword = {v: k for k, v in wordtoix.items()}
|
208 |
+
|
209 |
+
# Create fallback models
|
210 |
+
try:
|
211 |
+
text_encoder = RNN_ENCODER(len(wordtoix), nhidden=config['cond_dim']).to(device)
|
212 |
+
netG = NetG(config['nf'], config['z_dim'], config['cond_dim'], config['imsize'], config['ch_size']).to(device)
|
213 |
+
models_loaded = False
|
214 |
+
except Exception as e2:
|
215 |
+
print(f"Failed to create fallback models: {e2}")
|
216 |
+
|
217 |
+
return wordtoix, ixtoword, text_encoder, netG
|
218 |
|
219 |
+
# Try to load the models
|
220 |
+
try:
|
221 |
+
wordtoix, ixtoword, text_encoder, netG = load_models()
|
222 |
+
except Exception as e:
|
223 |
+
print(f"Error during model loading: {e}")
|
224 |
+
print("The application will attempt to continue but may not function correctly.")
|
225 |
|
226 |
def generate_image(text_input, num_images=1, seed=None):
|
227 |
"""Generate images from text description"""
|
228 |
if not text_input.strip():
|
229 |
+
return [Image.new('RGB', (256, 256), color='lightgray')] * num_images
|
|
|
|
|
|
|
|
|
|
|
230 |
|
231 |
+
try:
|
232 |
+
cap_array, cap_len = tokenize_and_build_captions(text_input, wordtoix)
|
233 |
+
|
234 |
+
if cap_len == 0:
|
235 |
+
return [Image.new('RGB', (256, 256), color='red')] * num_images
|
236 |
+
|
237 |
+
sent_emb = encode_caption(cap_array, cap_len, text_encoder, device)
|
238 |
+
|
239 |
+
# Set random seed if provided
|
240 |
+
if seed is not None:
|
241 |
+
random.seed(seed)
|
242 |
+
np.random.seed(seed)
|
243 |
+
torch.manual_seed(seed)
|
244 |
+
if torch.cuda.is_available():
|
245 |
+
torch.cuda.manual_seed_all(seed)
|
246 |
+
|
247 |
+
# Generate multiple images if requested
|
248 |
+
result_images = []
|
249 |
+
with torch.no_grad():
|
250 |
+
for _ in range(num_images):
|
251 |
+
# Generate noise
|
252 |
+
if config['truncation']:
|
253 |
+
noise = truncated_noise(1, config['z_dim'], config['trunc_rate'])
|
254 |
+
noise = torch.tensor(noise, dtype=torch.float).to(device)
|
255 |
+
else:
|
256 |
+
noise = torch.randn(1, config['z_dim']).to(device)
|
257 |
+
|
258 |
+
# Generate image
|
259 |
+
try:
|
260 |
+
fake_img = netG(noise, sent_emb)
|
261 |
+
img = save_img(fake_img[0])
|
262 |
+
result_images.append(img)
|
263 |
+
except Exception as e:
|
264 |
+
print(f"Error generating image: {e}")
|
265 |
+
# Return a placeholder image as fallback
|
266 |
+
img = Image.new('RGB', (256, 256), color=(255, 200, 200))
|
267 |
+
result_images.append(img)
|
268 |
+
|
269 |
+
return result_images
|
270 |
+
except Exception as e:
|
271 |
+
print(f"Error in generate_image: {e}")
|
272 |
+
traceback.print_exc()
|
273 |
+
return [Image.new('RGB', (256, 256), color='orange')] * num_images
|
274 |
+
|
275 |
+
# Create a simple message for model loading status
|
276 |
+
model_status = "✅ Models loaded successfully" if models_loaded_successfully else "⚠️ Using fallback models - images may not look good"
|
277 |
+
|
278 |
+
# Function to render error page if needed
|
279 |
+
def serve_error_page():
|
280 |
+
if os.path.exists('error_page.html'):
|
281 |
+
with open('error_page.html', 'r') as f:
|
282 |
+
return f.read()
|
283 |
+
else:
|
284 |
+
return "<html><body><h1>Error loading models</h1><p>The application failed to load the required models.</p></body></html>"
|
285 |
|
286 |
# Create Gradio interface
|
287 |
def generate_images_interface(text, num_images, random_seed):
|
288 |
+
seed = int(random_seed) if random_seed and random_seed.strip().isdigit() else None
|
289 |
return generate_image(text, num_images, seed)
|
290 |
|
291 |
+
# Create the Gradio interface
|
292 |
with gr.Blocks(title="Bird Image Generator") as demo:
|
293 |
+
if models_loaded_successfully:
|
294 |
+
# Normal interface when models loaded successfully
|
295 |
+
gr.Markdown("# Bird Image Generator using DF-GAN")
|
296 |
+
gr.Markdown("Enter a description of a bird and the model will generate corresponding images.")
|
297 |
+
|
298 |
+
gr.Markdown(f"**Model Status:** {model_status}")
|
299 |
+
|
300 |
+
with gr.Row():
|
301 |
+
with gr.Column():
|
302 |
+
text_input = gr.Textbox(
|
303 |
+
label="Bird Description",
|
304 |
+
placeholder="Enter a description of a bird (e.g., 'a small bird with a red head and black wings')",
|
305 |
+
lines=3
|
306 |
+
)
|
307 |
+
num_images = gr.Slider(minimum=1, maximum=4, value=1, step=1, label="Number of Images")
|
308 |
+
seed = gr.Textbox(label="Random Seed (optional)", placeholder="Leave empty for random results")
|
309 |
+
submit_btn = gr.Button("Generate Image")
|
310 |
+
|
311 |
+
with gr.Column():
|
312 |
+
image_output = gr.Gallery(label="Generated Images").style(grid=2, height="auto")
|
313 |
+
|
314 |
+
submit_btn.click(
|
315 |
+
fn=generate_images_interface,
|
316 |
+
inputs=[text_input, num_images, seed],
|
317 |
+
outputs=image_output
|
318 |
+
)
|
319 |
+
|
320 |
+
gr.Markdown("## Example Descriptions")
|
321 |
+
example_descriptions = [
|
322 |
+
"this bird has an orange bill, a white belly and white eyebrows",
|
323 |
+
"a small bird with a red head, breast, and belly and black wings",
|
324 |
+
"this bird is yellow with black and has a long, pointy beak",
|
325 |
+
"this bird is white in color, and has a orange beak"
|
326 |
+
]
|
327 |
+
|
328 |
+
gr.Examples(
|
329 |
+
examples=[[desc, 1, ""] for desc in example_descriptions],
|
330 |
+
inputs=[text_input, num_images, seed],
|
331 |
+
outputs=image_output,
|
332 |
+
fn=generate_images_interface
|
333 |
+
)
|
334 |
+
else:
|
335 |
+
# Modified interface with warning when models failed to load
|
336 |
+
gr.Markdown("# ⚠️ Bird Image Generator - Limited Functionality")
|
337 |
+
gr.Markdown("The pre-trained models could not be loaded correctly. The application will run with randomly initialized models.")
|
338 |
+
|
339 |
+
with gr.Row():
|
340 |
+
with gr.Column():
|
341 |
+
text_input = gr.Textbox(
|
342 |
+
label="Bird Description",
|
343 |
+
placeholder="Enter a description of a bird (e.g., 'a small bird with a red head and black wings')",
|
344 |
+
lines=3
|
345 |
+
)
|
346 |
+
num_images = gr.Slider(minimum=1, maximum=4, value=1, step=1, label="Number of Images")
|
347 |
+
seed = gr.Textbox(label="Random Seed (optional)", placeholder="Leave empty for random results")
|
348 |
+
submit_btn = gr.Button("Generate Image (Results will be random shapes)")
|
349 |
+
|
350 |
+
with gr.Column():
|
351 |
+
image_output = gr.Gallery(label="Generated Images (Random)").style(grid=2, height="auto")
|
352 |
+
|
353 |
+
submit_btn.click(
|
354 |
+
fn=generate_images_interface,
|
355 |
+
inputs=[text_input, num_images, seed],
|
356 |
+
outputs=image_output
|
357 |
+
)
|
358 |
+
|
359 |
+
gr.Markdown("""
|
360 |
+
### Model Loading Error
|
361 |
+
|
362 |
+
The application encountered an error while loading the pre-trained models. This could be due to:
|
363 |
+
|
364 |
+
1. Network connectivity issues
|
365 |
+
2. The model hosting service might be temporarily unavailable
|
366 |
+
3. The model files might have been moved or deleted
|
367 |
+
|
368 |
+
Please try refreshing the page or contact the Space owner if the issue persists.
|
369 |
+
""")
|
370 |
|
371 |
# Launch the app with appropriate configurations for Hugging Face Spaces
|
372 |
if __name__ == "__main__":
|
373 |
+
# Wait a moment before starting to make sure all logs are printed
|
374 |
+
time.sleep(1)
|
375 |
+
|
376 |
demo.launch(
|
377 |
server_name="0.0.0.0", # Bind to all network interfaces
|
378 |
share=False, # Don't use share links
|
download_models.py
CHANGED
@@ -1,10 +1,13 @@
|
|
1 |
import os
|
2 |
import sys
|
3 |
import subprocess
|
4 |
-
import gdown
|
5 |
import shutil
|
6 |
import nltk
|
7 |
from pathlib import Path
|
|
|
|
|
|
|
|
|
8 |
|
9 |
# Install NLTK data
|
10 |
nltk.download('punkt')
|
@@ -27,30 +30,91 @@ if not os.path.exists('DF-GAN/.git'):
|
|
27 |
|
28 |
print("Repository cloned and organized.")
|
29 |
|
30 |
-
#
|
31 |
-
|
32 |
-
|
33 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
34 |
|
35 |
-
#
|
36 |
-
|
37 |
-
|
|
|
38 |
|
39 |
-
#
|
40 |
-
|
|
|
41 |
captions_pickle_path = 'data/captions_DAMSM.pickle'
|
42 |
|
43 |
-
# Download
|
44 |
if not os.path.exists(bird_model_path):
|
45 |
print(f"Downloading bird model to {bird_model_path}...")
|
46 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
47 |
|
|
|
48 |
if not os.path.exists(text_encoder_path):
|
49 |
print(f"Downloading text encoder to {text_encoder_path}...")
|
50 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
51 |
|
|
|
52 |
if not os.path.exists(captions_pickle_path):
|
53 |
print(f"Downloading captions pickle to {captions_pickle_path}...")
|
54 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
55 |
|
56 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
import os
|
2 |
import sys
|
3 |
import subprocess
|
|
|
4 |
import shutil
|
5 |
import nltk
|
6 |
from pathlib import Path
|
7 |
+
import urllib.request
|
8 |
+
import zipfile
|
9 |
+
import torch
|
10 |
+
import time
|
11 |
|
12 |
# Install NLTK data
|
13 |
nltk.download('punkt')
|
|
|
30 |
|
31 |
print("Repository cloned and organized.")
|
32 |
|
33 |
+
# Function to download files with retries
|
34 |
+
def download_file(url, dest_path, max_retries=3):
|
35 |
+
for attempt in range(max_retries):
|
36 |
+
try:
|
37 |
+
print(f"Downloading from {url} to {dest_path} (attempt {attempt+1})")
|
38 |
+
urllib.request.urlretrieve(url, dest_path)
|
39 |
+
print(f"Successfully downloaded {dest_path}")
|
40 |
+
return True
|
41 |
+
except Exception as e:
|
42 |
+
print(f"Download attempt {attempt+1} failed: {e}")
|
43 |
+
time.sleep(2) # Wait before retrying
|
44 |
+
return False
|
45 |
|
46 |
+
# Model URLs - Changed to direct download URLs that are more reliable
|
47 |
+
BIRD_MODEL_URL = "https://huggingface.co/spaces/sayakpaul/df-gan-bird/resolve/main/state_epoch_1220.pth"
|
48 |
+
TEXT_ENCODER_URL = "https://huggingface.co/spaces/sayakpaul/df-gan-bird/resolve/main/text_encoder200.pth"
|
49 |
+
CAPTIONS_URL = "https://huggingface.co/spaces/sayakpaul/df-gan-bird/resolve/main/captions_DAMSM.pickle"
|
50 |
|
51 |
+
# Download paths
|
52 |
+
bird_model_path = 'data/state_epoch_1220.pth'
|
53 |
+
text_encoder_path = 'data/text_encoder200.pth'
|
54 |
captions_pickle_path = 'data/captions_DAMSM.pickle'
|
55 |
|
56 |
+
# Download bird model
|
57 |
if not os.path.exists(bird_model_path):
|
58 |
print(f"Downloading bird model to {bird_model_path}...")
|
59 |
+
success = download_file(BIRD_MODEL_URL, bird_model_path)
|
60 |
+
if not success:
|
61 |
+
print("Failed to download bird model after multiple attempts")
|
62 |
+
# Create a dummy model as fallback if needed
|
63 |
+
if not os.path.exists(bird_model_path):
|
64 |
+
print("Creating a dummy model for testing purposes...")
|
65 |
+
dummy_state = {
|
66 |
+
'model': {
|
67 |
+
'netG': {'dummy': torch.zeros(1)},
|
68 |
+
'netD': {'dummy': torch.zeros(1)},
|
69 |
+
'netC': {'dummy': torch.zeros(1)}
|
70 |
+
}
|
71 |
+
}
|
72 |
+
torch.save(dummy_state, bird_model_path)
|
73 |
+
print("Dummy model created as fallback")
|
74 |
|
75 |
+
# Download text encoder
|
76 |
if not os.path.exists(text_encoder_path):
|
77 |
print(f"Downloading text encoder to {text_encoder_path}...")
|
78 |
+
success = download_file(TEXT_ENCODER_URL, text_encoder_path)
|
79 |
+
if not success:
|
80 |
+
print("Failed to download text encoder after multiple attempts")
|
81 |
+
# Create a dummy encoder as fallback
|
82 |
+
if not os.path.exists(text_encoder_path):
|
83 |
+
print("Creating a dummy text encoder for testing purposes...")
|
84 |
+
dummy_encoder = {'dummy': torch.zeros(1)}
|
85 |
+
torch.save(dummy_encoder, text_encoder_path)
|
86 |
+
print("Dummy text encoder created as fallback")
|
87 |
|
88 |
+
# Download captions pickle
|
89 |
if not os.path.exists(captions_pickle_path):
|
90 |
print(f"Downloading captions pickle to {captions_pickle_path}...")
|
91 |
+
success = download_file(CAPTIONS_URL, captions_pickle_path)
|
92 |
+
if not success:
|
93 |
+
print("Failed to download captions pickle after multiple attempts")
|
94 |
+
# Create a placeholder pickle file for testing
|
95 |
+
if not os.path.exists(captions_pickle_path):
|
96 |
+
print("Creating a placeholder captions file...")
|
97 |
+
import pickle
|
98 |
+
wordtoix = {"the": 1, "bird": 2, "is": 3, "a": 4, "with": 5, "and": 6, "red": 7, "black": 8, "yellow": 9}
|
99 |
+
ixtoword = {v: k for k, v in wordtoix.items()}
|
100 |
+
test_data = [None, None, ixtoword, wordtoix]
|
101 |
+
with open(captions_pickle_path, 'wb') as f:
|
102 |
+
pickle.dump(test_data, f)
|
103 |
+
print("Placeholder captions file created as fallback")
|
104 |
+
|
105 |
+
# Verify downloads
|
106 |
+
all_files_exist = (
|
107 |
+
os.path.exists(bird_model_path) and
|
108 |
+
os.path.exists(text_encoder_path) and
|
109 |
+
os.path.exists(captions_pickle_path)
|
110 |
+
)
|
111 |
|
112 |
+
if all_files_exist:
|
113 |
+
print("All model files downloaded and prepared successfully!")
|
114 |
+
else:
|
115 |
+
missing_files = []
|
116 |
+
if not os.path.exists(bird_model_path): missing_files.append(bird_model_path)
|
117 |
+
if not os.path.exists(text_encoder_path): missing_files.append(text_encoder_path)
|
118 |
+
if not os.path.exists(captions_pickle_path): missing_files.append(captions_pickle_path)
|
119 |
+
print(f"Warning: The following files could not be downloaded: {', '.join(missing_files)}")
|
120 |
+
print("The application may not function correctly.")
|
error_page.html
ADDED
@@ -0,0 +1,99 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
<!DOCTYPE html>
|
2 |
+
<html>
|
3 |
+
<head>
|
4 |
+
<title>DF-GAN Bird Generator - Model Loading Issue</title>
|
5 |
+
<style>
|
6 |
+
body {
|
7 |
+
font-family: Arial, sans-serif;
|
8 |
+
line-height: 1.6;
|
9 |
+
margin: 0;
|
10 |
+
padding: 20px;
|
11 |
+
background-color: #f8f9fa;
|
12 |
+
color: #333;
|
13 |
+
}
|
14 |
+
.container {
|
15 |
+
max-width: 800px;
|
16 |
+
margin: 40px auto;
|
17 |
+
padding: 30px;
|
18 |
+
background: white;
|
19 |
+
border-radius: 10px;
|
20 |
+
box-shadow: 0 0 20px rgba(0,0,0,0.1);
|
21 |
+
}
|
22 |
+
h1 {
|
23 |
+
color: #d9534f;
|
24 |
+
margin-bottom: 20px;
|
25 |
+
}
|
26 |
+
h2 {
|
27 |
+
color: #333;
|
28 |
+
margin-top: 30px;
|
29 |
+
}
|
30 |
+
pre {
|
31 |
+
background-color: #f5f5f5;
|
32 |
+
padding: 15px;
|
33 |
+
border-radius: 5px;
|
34 |
+
overflow-x: auto;
|
35 |
+
}
|
36 |
+
.warning {
|
37 |
+
background-color: #fff3cd;
|
38 |
+
border-left: 5px solid #ffc107;
|
39 |
+
padding: 15px;
|
40 |
+
margin: 20px 0;
|
41 |
+
border-radius: 5px;
|
42 |
+
}
|
43 |
+
.error {
|
44 |
+
background-color: #f8d7da;
|
45 |
+
border-left: 5px solid #dc3545;
|
46 |
+
padding: 15px;
|
47 |
+
margin: 20px 0;
|
48 |
+
border-radius: 5px;
|
49 |
+
}
|
50 |
+
.success {
|
51 |
+
background-color: #d4edda;
|
52 |
+
border-left: 5px solid #28a745;
|
53 |
+
padding: 15px;
|
54 |
+
margin: 20px 0;
|
55 |
+
border-radius: 5px;
|
56 |
+
}
|
57 |
+
</style>
|
58 |
+
</head>
|
59 |
+
<body>
|
60 |
+
<div class="container">
|
61 |
+
<h1>DF-GAN Bird Generator - Model Loading Issue</h1>
|
62 |
+
|
63 |
+
<div class="error">
|
64 |
+
<p><strong>There was an issue loading the required model files.</strong></p>
|
65 |
+
<p>The application is running in fallback mode with randomly initialized weights. Generated images will not look like realistic birds.</p>
|
66 |
+
</div>
|
67 |
+
|
68 |
+
<h2>What happened?</h2>
|
69 |
+
<p>The application tried to download the pre-trained DF-GAN model files but encountered an error. This could be due to:</p>
|
70 |
+
<ul>
|
71 |
+
<li>Network connectivity issues</li>
|
72 |
+
<li>The model hosting service might be temporarily unavailable</li>
|
73 |
+
<li>The model files might have been moved or deleted</li>
|
74 |
+
</ul>
|
75 |
+
|
76 |
+
<h2>What can you do?</h2>
|
77 |
+
<p>Here are some options to fix this issue:</p>
|
78 |
+
<ol>
|
79 |
+
<li>Refresh the page and try again - the issue might be temporary</li>
|
80 |
+
<li>Contact the Space owner to notify them of the issue</li>
|
81 |
+
<li>If you're the owner, check that the model files are correctly hosted</li>
|
82 |
+
</ol>
|
83 |
+
|
84 |
+
<div class="success">
|
85 |
+
<p>The application will still run, but with reduced functionality. You can still enter text descriptions, but the generated images will not be realistic.</p>
|
86 |
+
</div>
|
87 |
+
|
88 |
+
<h2>Technical Details</h2>
|
89 |
+
<p>The application was unable to download or load one or more of the following files:</p>
|
90 |
+
<ul>
|
91 |
+
<li>state_epoch_1220.pth (Generator model)</li>
|
92 |
+
<li>text_encoder200.pth (Text encoder model)</li>
|
93 |
+
<li>captions_DAMSM.pickle (Vocabulary data)</li>
|
94 |
+
</ul>
|
95 |
+
|
96 |
+
<p>Check the application logs for more detailed error information.</p>
|
97 |
+
</div>
|
98 |
+
</body>
|
99 |
+
</html>
|
nltk_setup.py
ADDED
@@ -0,0 +1,16 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import nltk
|
2 |
+
import os
|
3 |
+
|
4 |
+
# Make sure NLTK data directory exists
|
5 |
+
nltk_data_dir = os.path.expanduser('~/nltk_data')
|
6 |
+
os.makedirs(nltk_data_dir, exist_ok=True)
|
7 |
+
|
8 |
+
# Check if punkt tokenizer already exists
|
9 |
+
punkt_dir = os.path.join(nltk_data_dir, 'tokenizers', 'punkt')
|
10 |
+
if not os.path.exists(punkt_dir):
|
11 |
+
print("Downloading NLTK punkt tokenizer...")
|
12 |
+
nltk.download('punkt', quiet=False)
|
13 |
+
else:
|
14 |
+
print("NLTK punkt tokenizer already exists")
|
15 |
+
|
16 |
+
print("NLTK setup complete")
|
requirements.txt
CHANGED
@@ -1,16 +1,10 @@
|
|
1 |
-
flask==2.0.1
|
2 |
torch>=1.9.0
|
3 |
torchvision>=0.10.0
|
4 |
Pillow>=9.0.0
|
5 |
-
nltk>=3.6.0
|
6 |
-
gunicorn==20.1.0
|
7 |
-
python-dotenv==0.19.0
|
8 |
-
requests==2.26.0
|
9 |
-
matplotlib==3.5.1
|
10 |
-
tqdm>=4.62.0
|
11 |
numpy>=1.20.0
|
|
|
|
|
12 |
scipy>=1.7.0
|
13 |
omegaconf>=2.1.0
|
14 |
gradio>=3.50.0
|
15 |
-
easydict>=1.9
|
16 |
-
gdown>=4.6.0
|
|
|
|
|
1 |
torch>=1.9.0
|
2 |
torchvision>=0.10.0
|
3 |
Pillow>=9.0.0
|
|
|
|
|
|
|
|
|
|
|
|
|
4 |
numpy>=1.20.0
|
5 |
+
tqdm>=4.62.0
|
6 |
+
nltk>=3.6.0
|
7 |
scipy>=1.7.0
|
8 |
omegaconf>=2.1.0
|
9 |
gradio>=3.50.0
|
10 |
+
easydict>=1.9
|
|
startup.sh
CHANGED
@@ -1,10 +1,19 @@
|
|
1 |
#!/bin/bash
|
|
|
|
|
|
|
2 |
|
3 |
# Install NLTK data
|
4 |
-
|
|
|
5 |
|
6 |
# Run the download_models.py script to get the models
|
7 |
-
|
|
|
|
|
|
|
|
|
8 |
|
9 |
# Start the Gradio app
|
10 |
-
|
|
|
|
1 |
#!/bin/bash
|
2 |
+
set -e
|
3 |
+
|
4 |
+
echo "Starting DF-GAN Bird Image Generator setup..."
|
5 |
|
6 |
# Install NLTK data
|
7 |
+
echo "Setting up NLTK data..."
|
8 |
+
python nltk_setup.py
|
9 |
|
10 |
# Run the download_models.py script to get the models
|
11 |
+
echo "Downloading model files..."
|
12 |
+
python download_models.py || {
|
13 |
+
echo "Warning: Some model files may not have downloaded correctly."
|
14 |
+
echo "The application will attempt to continue with fallback models."
|
15 |
+
}
|
16 |
|
17 |
# Start the Gradio app
|
18 |
+
echo "Starting the web application..."
|
19 |
+
exec python app.py
|