lisonallen commited on
Commit
ae398e9
·
1 Parent(s): 8c6d7f4

Switch to API-based image generation for better compatibility with Hugging Face Spaces

Browse files
Files changed (2) hide show
  1. app.py +114 -70
  2. requirements.txt +3 -8
app.py CHANGED
@@ -4,6 +4,10 @@ import random
4
  import logging
5
  import sys
6
  import os
 
 
 
 
7
  from PIL import Image as PILImage
8
 
9
  # 设置日志记录
@@ -75,74 +79,112 @@ def create_backup_image(prompt=""):
75
  # 预加载图像用于快速响应
76
  PLACEHOLDER_IMAGE = create_backup_image("placeholder")
77
 
78
- # 尝试导入必要的AI库
79
- try:
80
- import torch
81
- from diffusers import StableDiffusionPipeline
82
-
83
- HAS_AI_LIBS = True
84
- logger.info("Successfully imported AI libraries")
85
- except ImportError as e:
86
- logger.error(f"Failed to import AI libraries: {e}")
87
- HAS_AI_LIBS = False
88
-
89
- # AI 模型加载和图像生成
90
- def generate_ai_image(prompt, seed=None):
91
- if not HAS_AI_LIBS:
92
- logger.error("AI libraries not available")
93
- return PLACEHOLDER_IMAGE
 
 
94
 
95
- # 设置随机种子
96
- if seed is None:
97
- seed = random.randint(0, 2147483647)
 
 
 
 
 
98
 
99
  try:
100
- logger.info(f"Generating image for: {prompt}")
101
-
102
- # 使用兼容的旧版本API加载模型
103
- model_id = "runwayml/stable-diffusion-v1-5"
104
- logger.info(f"Loading model: {model_id}")
105
-
106
- device = "cuda" if torch.cuda.is_available() else "cpu"
107
- torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
 
 
108
 
109
- # 加载模型 - 使用兼容的低级API
110
- pipe = StableDiffusionPipeline.from_pretrained(
111
- model_id,
112
- torch_dtype=torch_dtype,
113
- use_auth_token=False, # 明确不使用认证
114
- safety_checker=None, # 禁用安全检查器
115
- )
116
- pipe = pipe.to(device)
117
-
118
- # 优化内存
119
- if torch.cuda.is_available():
120
- pipe.enable_attention_slicing()
121
- torch.cuda.empty_cache()
122
 
123
- logger.info("Model loaded, generating image...")
 
 
 
 
 
 
 
 
 
 
124
 
125
- # 生成图像
126
- generator = torch.Generator(device).manual_seed(seed)
127
- image = pipe(
128
- prompt=prompt,
129
- guidance_scale=7.5,
130
- num_inference_steps=4, # 最小步数
131
- generator=generator,
132
- height=512,
133
- width=512
134
- ).images[0]
 
 
 
 
 
 
135
 
136
- # 清理缓存
137
- if torch.cuda.is_available():
138
- torch.cuda.empty_cache()
139
-
140
- logger.info(f"Image generation successful with seed: {seed}")
141
- return image
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
142
 
 
 
 
143
  except Exception as e:
144
- logger.error(f"AI image generation failed: {e}")
145
- return create_backup_image(prompt)
146
 
147
  # 使用简单的规则生成图像作为备用方案
148
  def generate_rule_based_image(prompt):
@@ -235,17 +277,19 @@ def generate_image(prompt):
235
 
236
  logger.info(f"Received prompt: {prompt}")
237
 
238
- # 尝试使用AI生成
239
- if HAS_AI_LIBS:
240
- try:
241
- image = generate_ai_image(prompt)
242
- if image is not None:
243
- return image
244
- except Exception as e:
245
- logger.error(f"Error using AI generation: {e}")
 
 
246
 
247
- # 如果AI不可用或失败,使用规则生成
248
- logger.warning("Using rule-based image generation")
249
  return generate_rule_based_image(prompt)
250
 
251
  # 为旧版 gradio 创建界面
@@ -260,7 +304,7 @@ def create_demo():
260
  ),
261
  outputs=gr.Image(label="Generated Image", type="pil"),
262
  title="Text to Image Generator",
263
- description="Enter a text description to generate an image.",
264
  examples=[
265
  "a cute cat sitting on a windowsill",
266
  "beautiful sunset over mountains",
 
4
  import logging
5
  import sys
6
  import os
7
+ import requests
8
+ import io
9
+ import json
10
+ import base64
11
  from PIL import Image as PILImage
12
 
13
  # 设置日志记录
 
79
  # 预加载图像用于快速响应
80
  PLACEHOLDER_IMAGE = create_backup_image("placeholder")
81
 
82
+ # 使用 Hugging Face Inference API 生成图像
83
+ def generate_image_with_api(prompt, api_url=None, api_key=None):
84
+ """
85
+ 使用 Hugging Face Inference API 生成图像
86
+
87
+ Parameters:
88
+ - prompt: 文本提示
89
+ - api_url: API 端点 URL (可选)
90
+ - api_key: API 密钥 (可选)
91
+
92
+ Returns:
93
+ - PIL Image
94
+ """
95
+ logger.info(f"Generating image via API for: {prompt}")
96
+
97
+ # 默认使用 Stable Diffusion API
98
+ if api_url is None:
99
+ api_url = "https://api-inference.huggingface.co/models/runwayml/stable-diffusion-v1-5"
100
 
101
+ # 尝试从环境变量中获取 API 密钥
102
+ if api_key is None:
103
+ api_key = os.environ.get("HF_API_KEY", "")
104
+
105
+ # 如果没有 API 密钥,使用公共访问(可能会受到速率限制)
106
+ headers = {}
107
+ if api_key:
108
+ headers["Authorization"] = f"Bearer {api_key}"
109
 
110
  try:
111
+ # 设置请求参数
112
+ payload = {
113
+ "inputs": prompt,
114
+ "parameters": {
115
+ "num_inference_steps": 10, # 减少推理步骤以加快速度
116
+ "guidance_scale": 7.5,
117
+ "width": 512,
118
+ "height": 512
119
+ }
120
+ }
121
 
122
+ # 发送请求
123
+ response = requests.post(api_url, headers=headers, json=payload)
 
 
 
 
 
 
 
 
 
 
 
124
 
125
+ # 检查响应是否成功
126
+ if response.status_code == 200:
127
+ # 从响应中获取图像
128
+ image = PILImage.open(io.BytesIO(response.content))
129
+ logger.info("Successfully generated image via API")
130
+ return image
131
+ else:
132
+ # 如果遇到错误,记录响应并返回备用图像
133
+ error_text = response.text
134
+ logger.error(f"API error: {response.status_code}, {error_text}")
135
+ return None
136
 
137
+ except Exception as e:
138
+ logger.error(f"Failed to generate image via API: {e}")
139
+ return None
140
+
141
+ # 使用 Hugging Face Spaces 内置模型生成图像
142
+ def generate_image_with_spaces(prompt):
143
+ """使用同一个 Hugging Face Space 内的其他公共空间来生成图像"""
144
+ logger.info(f"Generating image via Spaces for: {prompt}")
145
+
146
+ try:
147
+ # 一些公共可用的 Stable Diffusion 空间 URLs
148
+ space_urls = [
149
+ "https://huggingface-projects-stable-diffusion-demo.hf.space/api/predict",
150
+ "https://huggingface-projects-text-to-image.hf.space/api/predict",
151
+ "https://dataautogpt-playground.hf.space/api/predict"
152
+ ]
153
 
154
+ # 尝试每个 URL 直到成功
155
+ for url in space_urls:
156
+ try:
157
+ payload = {
158
+ "data": [prompt, 7.5, 512, 512]
159
+ }
160
+ response = requests.post(url, json=payload, timeout=30)
161
+
162
+ if response.status_code == 200:
163
+ # 解析 JSON 响应
164
+ result = response.json()
165
+
166
+ # 根据结果格式提取图像数据
167
+ if isinstance(result, dict) and 'data' in result:
168
+ image_data = result['data'][0]
169
+
170
+ # 检查图像数据格式
171
+ if image_data.startswith('data:image'):
172
+ # 提取 base64 图像数据
173
+ image_b64 = image_data.split(',')[1]
174
+ image_bytes = base64.b64decode(image_b64)
175
+ image = PILImage.open(io.BytesIO(image_bytes))
176
+ logger.info(f"Successfully generated image via {url}")
177
+ return image
178
+ except Exception as e:
179
+ logger.warning(f"Failed to generate with {url}: {e}")
180
+ continue
181
 
182
+ logger.error("All space URLs failed")
183
+ return None
184
+
185
  except Exception as e:
186
+ logger.error(f"Failed to generate image via Spaces: {e}")
187
+ return None
188
 
189
  # 使用简单的规则生成图像作为备用方案
190
  def generate_rule_based_image(prompt):
 
277
 
278
  logger.info(f"Received prompt: {prompt}")
279
 
280
+ # 尝试使用 Hugging Face API 生成
281
+ image = generate_image_with_api(prompt)
282
+ if image is not None:
283
+ return image
284
+
285
+ # 如果 API 方法失败,尝试使用 Spaces
286
+ logger.info("API method failed, trying Spaces method...")
287
+ image = generate_image_with_spaces(prompt)
288
+ if image is not None:
289
+ return image
290
 
291
+ # 如果所有 AI 方法都失败,使用规则生成
292
+ logger.warning("All AI methods failed, using rule-based image generation")
293
  return generate_rule_based_image(prompt)
294
 
295
  # 为旧版 gradio 创建界面
 
304
  ),
305
  outputs=gr.Image(label="Generated Image", type="pil"),
306
  title="Text to Image Generator",
307
+ description="Enter a text description to generate an image. This demo uses Hugging Face API for image generation.",
308
  examples=[
309
  "a cute cat sitting on a windowsill",
310
  "beautiful sunset over mountains",
requirements.txt CHANGED
@@ -1,8 +1,3 @@
1
- accelerate>=0.15.0
2
- diffusers==0.10.2
3
- huggingface-hub>=0.13.0
4
- torch==1.13.1
5
- transformers>=4.25.1
6
- safetensors>=0.3.1
7
- gradio>=3.19.1,<3.25.0
8
- Pillow>=9.0.0,<10.0.0
 
1
+ requests>=2.26.0
2
+ Pillow>=9.0.0,<10.0.0
3
+ gradio>=3.19.1,<3.25.0