Update app.py

app.py

@@ -2,27 +2,15 @@ import os
 
 import gradio as gr
 import torch
-from diffusers import StableDiffusionPipeline
 from PIL import Image
-
-import ImageReward as RM
-
-# initialize
-
-model_id = "runwayml/stable-diffusion-v1-5"
-pipe = StableDiffusionPipeline.from_pretrained(
-    model_id,
-    torch_dtype=torch.float16,
-)
-
-model = RM.load("ImageReward-v1.0")
+import io
+import base64
+import requests
+import json
 
 images_in_gallery = []
 rewards_in_gallery = []
 
-# event functions
-
-
 def generate_images(
     prompt, magic_words, num, height, width, num_inference_steps, guidance_scale
 ):
@@ -31,14 +19,33 @@ def generate_images(
     if magic_words is not None:
         prompt += ", ".join(magic_words)
 
-    images_in_gallery = pipe(
-        prompt,
-        height=height,
-        width=width,
-        num_inference_steps=num_inference_steps,
-        guidance_scale=guidance_scale,
-        num_images_per_prompt=num,
-    ).images
+    # post 请求发送到服务器
+
+    # 定义请求的 URL 和数据
+    url = 'https://tianqi.aminer.cn/image_reward_hf/generate_image'
+    data = {'prompt': prompt, 
+            'height': height, 
+            'width':width, 
+            'num_inference_steps':num_inference_steps,
+            'guidance_scale':guidance_scale,
+            'num':num
+            }
+    headers = {'Content-Type': 'application/json'}
+
+    # 发送 POST 请求
+    data = json.dumps(data)
+    response = requests.post(url, data=data, headers=headers)
+    image_ls = response.json()['image_list']
+    
+    images_in_gallery = []
+    for base_image in image_ls:
+        image_bytes = base64.b64decode(base_image)
+        # 创建 BytesIO 对象并读取图像字节流
+        image_stream = io.BytesIO(image_bytes)
+        # 打开图像
+        image = Image.open(image_stream)
+        images_in_gallery.append(image)
+
     rewards_in_gallery = [None] * len(images_in_gallery)
     return list(zip(images_in_gallery, rewards_in_gallery))
 
@@ -50,9 +57,25 @@ def score_and_rank(prompt):
 
     if num_not_scored > 0:
         images_to_score = images_in_gallery[-num_not_scored:]
+        image_ls = []
+        for image in images_to_score:
+            image_bytes = io.BytesIO()
+            image.save(image_bytes, format='JPEG')
+            image_bytes.seek(0)
+            # 将字节流转换为 Base64 编码
+            base64_image = base64.b64encode(image_bytes.read()).decode('utf-8')
+            image_ls.append(base64_image)
         with torch.no_grad():
-            ranks, rewards = model.inference_rank(prompt, images_to_score)
-
+            # post 请求发送到服务器
+            url = 'https://tianqi.aminer.cn/image_reward_hf/score_and_rank'
+            data = {'images_to_score': image_ls, 'prompt':prompt}
+            data = json.dumps(data)
+            headers = {'Content-Type': 'application/json'}
+
+            # 发送 POST 请求
+            response = requests.post(url, data=data, headers=headers)
+            rewards = response.json()['rewards']
+            
         if not isinstance(rewards, list):
             rewards = [rewards]
         rewards_in_gallery = rewards_in_gallery[:-num_not_scored] + rewards