Update image generation to use InferenceClient and adjust requirements

app.py

@@ -19,7 +19,7 @@ from diffusers import FluxPipeline
 from huggingface_hub import InferenceClient
 
 llm_client = Client("Qwen/Qwen2.5-72B-Instruct")
-client = InferenceClient("black-forest-labs/FLUX.1-dev")
+t2i_client = Client("black-forest-labs/FLUX.1-dev")
 
 # device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 # pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16).to("cpu")
@@ -54,18 +54,6 @@ Focus on the item itself, ensuring it is fully described, and specify a plain, w
 
     return object_t2i_prompt
 
-def generate_item_image(object_t2i_prompt):
-    # image = pipe(prompt=object_t2i_prompt, guidance_scale=3.5, num_inference_steps=28, width=1024, height=1024, generator=torch.Generator("cpu").manual_seed(0), output_type="pil").images[0]
-    image = client.text_to_image(object_t2i_prompt, guidance_scale=3.5, num_inference_steps=28, width=1024, height=1024)
-    trial_id, processed_image = preprocess_pil_image(image)
-    return trial_id, processed_image
-
-MAX_SEED = np.iinfo(np.int32).max
-TMP_DIR = "/tmp/Trellis-demo"
-
-os.makedirs(TMP_DIR, exist_ok=True)
-
-
 def preprocess_pil_image(image: Image.Image) -> Tuple[str, Image.Image]:
     """
     Preprocess the input image.
@@ -82,6 +70,27 @@ def preprocess_pil_image(image: Image.Image) -> Tuple[str, Image.Image]:
     processed_image.save(f"{TMP_DIR}/{trial_id}.png")
     return trial_id, processed_image
 
+def generate_item_image(object_t2i_prompt):
+    # image = pipe(prompt=object_t2i_prompt, guidance_scale=3.5, num_inference_steps=28, width=1024, height=1024, generator=torch.Generator("cpu").manual_seed(0), output_type="pil").images[0]
+    # image = client.text_to_image(object_t2i_prompt, guidance_scale=3.5, num_inference_steps=28, width=1024, height=1024)
+    img_path = t2i_client.predict(
+		prompt=object_t2i_prompt,
+		seed=0,
+		randomize_seed=True,
+		width=1024,
+		height=1024,
+		guidance_scale=3.5,
+		num_inference_steps=8,
+		api_name="/infer"
+    )[0]
+    image = Image.open(img_path)
+    trial_id, processed_image = preprocess_pil_image(image)
+    return trial_id, processed_image
+
+MAX_SEED = np.iinfo(np.int32).max
+TMP_DIR = "/tmp/Trellis-demo"
+
+os.makedirs(TMP_DIR, exist_ok=True)
 
 def pack_state(gs: Gaussian, mesh: MeshExtractResult, trial_id: str) -> dict:
     return {

requirements.txt

@@ -4,7 +4,7 @@
 accelerate
 sentencepiece
 diffusers
-gradio_client
+gradio_client==1.4.0
 huggingface-hub==0.26.5
 torch==2.4.0
 torchvision==0.19.0