akhaliq HF Staff commited on
Commit
fc5bd53
·
verified ·
1 Parent(s): 17cc4e0

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +133 -4
app.py CHANGED
@@ -24,7 +24,7 @@ def get_headers():
24
  "X-HF-Bill-To": "huggingface"
25
  }
26
 
27
- def query_api(payload):
28
  """Send request to the API and return response"""
29
  headers = get_headers()
30
 
@@ -33,12 +33,141 @@ def query_api(payload):
33
  payload["inputs"] = base64.b64encode(payload["image_bytes"]).decode("utf-8")
34
  del payload["image_bytes"]
35
 
 
36
  response = requests.post(API_URL, headers=headers, json=payload)
37
 
38
  if response.status_code != 200:
39
  raise gr.Error(f"API request failed with status {response.status_code}: {response.text}")
40
 
41
- return response.content
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
42
 
43
  def upload_image_to_fal(image_bytes):
44
  """Upload image to fal.ai and return the URL"""
@@ -227,8 +356,8 @@ def chat_fn(message, chat_history, seed, randomize_seed, guidance_scale, steps,
227
  progress(0.1, desc="Generating image...")
228
 
229
  try:
230
- # Make API request
231
- image_bytes = query_api(payload)
232
 
233
  # Try to convert response bytes to PIL Image
234
  try:
 
24
  "X-HF-Bill-To": "huggingface"
25
  }
26
 
27
+ def query_api(payload, progress_callback=None):
28
  """Send request to the API and return response"""
29
  headers = get_headers()
30
 
 
33
  payload["inputs"] = base64.b64encode(payload["image_bytes"]).decode("utf-8")
34
  del payload["image_bytes"]
35
 
36
+ # Submit the job
37
  response = requests.post(API_URL, headers=headers, json=payload)
38
 
39
  if response.status_code != 200:
40
  raise gr.Error(f"API request failed with status {response.status_code}: {response.text}")
41
 
42
+ # Debug the response
43
+ print(f"Response status: {response.status_code}")
44
+ print(f"Response headers: {dict(response.headers)}")
45
+ print(f"Response content type: {response.headers.get('content-type', 'unknown')}")
46
+ print(f"Response content length: {len(response.content)}")
47
+ print(f"First 500 chars of response: {response.content[:500]}")
48
+
49
+ # Check if response is JSON (queue status) or binary (direct image)
50
+ content_type = response.headers.get('content-type', '').lower()
51
+
52
+ if 'application/json' in content_type:
53
+ # Response is JSON, likely queue status
54
+ try:
55
+ json_response = response.json()
56
+ print(f"JSON response: {json_response}")
57
+
58
+ # Check if job was queued
59
+ if json_response.get("status") == "IN_QUEUE":
60
+ request_id = json_response.get("request_id")
61
+ if not request_id:
62
+ raise gr.Error("No request_id provided in queue response")
63
+
64
+ # Poll for completion using the proper HF router endpoints
65
+ max_attempts = 60 # Wait up to 5 minutes
66
+ attempt = 0
67
+
68
+ while attempt < max_attempts:
69
+ if progress_callback:
70
+ progress_callback(0.1 + (attempt / max_attempts) * 0.8, f"Processing... (attempt {attempt + 1}/60)")
71
+
72
+ time.sleep(5) # Wait 5 seconds between polls
73
+
74
+ # Check status using HF router format
75
+ status_url = f"https://router.huggingface.co/fal-ai/fal-ai/flux-kontext/dev/requests/{request_id}/status"
76
+ status_response = requests.get(status_url, headers=headers)
77
+
78
+ if status_response.status_code != 200:
79
+ print(f"Status response: {status_response.status_code} - {status_response.text}")
80
+ # Continue polling even if status check fails temporarily
81
+ attempt += 1
82
+ continue
83
+
84
+ try:
85
+ status_data = status_response.json()
86
+ print(f"Status check {attempt + 1}: {status_data}")
87
+
88
+ if status_data.get("status") == "COMPLETED":
89
+ # Job completed, get the result
90
+ result_url = f"https://router.huggingface.co/fal-ai/fal-ai/flux-kontext/dev/requests/{request_id}"
91
+ result_response = requests.get(result_url, headers=headers)
92
+
93
+ if result_response.status_code != 200:
94
+ print(f"Result response: {result_response.status_code} - {result_response.text}")
95
+ raise gr.Error(f"Failed to get result: {result_response.status_code}")
96
+
97
+ # Check if result is direct image bytes or JSON
98
+ result_content_type = result_response.headers.get('content-type', '').lower()
99
+ if 'image/' in result_content_type:
100
+ # Direct image bytes
101
+ return result_response.content
102
+ else:
103
+ # Try to parse as JSON for image URL or base64
104
+ try:
105
+ result_data = result_response.json()
106
+ print(f"Result data: {result_data}")
107
+
108
+ # Look for images in various formats
109
+ if 'images' in result_data and len(result_data['images']) > 0:
110
+ image_info = result_data['images'][0]
111
+ if isinstance(image_info, dict) and 'url' in image_info:
112
+ # Download the image
113
+ img_response = requests.get(image_info['url'])
114
+ return img_response.content
115
+ elif isinstance(image_info, str):
116
+ # Base64 encoded
117
+ return base64.b64decode(image_info)
118
+ elif 'image' in result_data:
119
+ # Single image field
120
+ if isinstance(result_data['image'], str):
121
+ return base64.b64decode(result_data['image'])
122
+ else:
123
+ # Maybe it's direct image bytes
124
+ return result_response.content
125
+
126
+ except requests.exceptions.JSONDecodeError:
127
+ # Result might be direct image bytes
128
+ return result_response.content
129
+
130
+ elif status_data.get("status") == "FAILED":
131
+ error_msg = status_data.get("error", "Unknown error")
132
+ raise gr.Error(f"Job failed: {error_msg}")
133
+
134
+ # Still processing, continue polling
135
+ attempt += 1
136
+
137
+ except requests.exceptions.JSONDecodeError:
138
+ print("Failed to parse status response, continuing...")
139
+ attempt += 1
140
+ continue
141
+
142
+ raise gr.Error("Job timed out after 5 minutes")
143
+
144
+ elif json_response.get("status") == "COMPLETED":
145
+ # Job completed immediately
146
+ if 'images' in json_response and len(json_response['images']) > 0:
147
+ image_info = json_response['images'][0]
148
+ if isinstance(image_info, dict) and 'url' in image_info:
149
+ img_response = requests.get(image_info['url'])
150
+ return img_response.content
151
+ elif isinstance(image_info, str):
152
+ return base64.b64decode(image_info)
153
+ elif 'image' in json_response:
154
+ return base64.b64decode(json_response['image'])
155
+ else:
156
+ raise gr.Error(f"No images found in immediate response: {json_response}")
157
+
158
+ else:
159
+ raise gr.Error(f"Unexpected response status: {json_response.get('status', 'unknown')}")
160
+
161
+ except requests.exceptions.JSONDecodeError as e:
162
+ raise gr.Error(f"Failed to parse JSON response: {str(e)}")
163
+
164
+ elif 'image/' in content_type:
165
+ # Response is direct image bytes
166
+ return response.content
167
+
168
+ else:
169
+ # Unknown content type, try to handle as bytes
170
+ return response.content
171
 
172
  def upload_image_to_fal(image_bytes):
173
  """Upload image to fal.ai and return the URL"""
 
356
  progress(0.1, desc="Generating image...")
357
 
358
  try:
359
+ # Make API request with progress callback
360
+ image_bytes = query_api(payload, progress_callback=progress)
361
 
362
  # Try to convert response bytes to PIL Image
363
  try: