Spaces:
Running
on
Zero
Running
on
Zero
Update app.py
Browse files
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 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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:
|