Update fal_recraft_upscale.py

This commit is contained in:
nofunstudio
2025-06-08 20:16:28 -06:00
parent 8d287d4b99
commit 76083b4331
+54 -17
View File
@@ -133,15 +133,34 @@ class FalAPI_recraft_upscale:
# 1) Set up FAL client with API token
os.environ["FAL_KEY"] = api_token
print(f"[Recraft Upscale] API token set, processing image...")
# Check if fal_client is working
try:
print(f"[Recraft Upscale] Testing fal_client import: {fal_client.__version__ if hasattr(fal_client, '__version__') else 'version unknown'}")
except Exception as e:
print(f"[Recraft Upscale] fal_client check failed: {str(e)}")
# 2) Convert ComfyUI image tensor to temporary file
image_file = self.tensor_to_tempfile(image)
print(f"[Recraft Upscale] Created temp file: {image_file.name}")
try:
temp_file_path = self.tensor_to_tempfile(image)
print(f"[Recraft Upscale] Created temp file: {temp_file_path}")
except Exception as e:
print(f"[Recraft Upscale] Error creating temp file: {str(e)}")
raise
# 3) Upload image to FAL storage and get URL
print(f"[Recraft Upscale] Uploading image to FAL storage...")
image_url = fal_client.upload_file(image_file)
print(f"[Recraft Upscale] Image uploaded successfully: {image_url}")
try:
print(f"[Recraft Upscale] Uploading image to FAL storage...")
image_url = fal_client.upload_file(temp_file_path)
print(f"[Recraft Upscale] Image uploaded successfully: {image_url}")
except Exception as e:
print(f"[Recraft Upscale] Error uploading file: {str(e)}")
# Clean up file before re-raising
try:
os.remove(temp_file_path)
except:
pass
raise
# 4) Build the input arguments
arguments = {
@@ -152,17 +171,25 @@ class FalAPI_recraft_upscale:
print(f"[Recraft Upscale] Starting upscale with safety_checker={enable_safety_checker}")
# 5) Call fal_client.subscribe()
result = fal_client.subscribe(
"fal-ai/recraft/upscale/crisp",
arguments=arguments,
with_logs=True,
on_queue_update=self.on_queue_update,
)
print(f"[Recraft Upscale] Upscale completed successfully")
try:
result = fal_client.subscribe(
"fal-ai/recraft/upscale/crisp",
arguments=arguments,
with_logs=True,
on_queue_update=self.on_queue_update,
)
print(f"[Recraft Upscale] Upscale completed successfully")
except Exception as e:
print(f"[Recraft Upscale] Error during upscale: {str(e)}")
# Clean up file before re-raising
try:
os.remove(temp_file_path)
except:
pass
raise
# 6) Clean up temporary file
image_file.close()
os.remove(image_file.name)
os.remove(temp_file_path)
print(f"[Recraft Upscale] Cleaned up temp file")
if not result or not result.get("image"):
@@ -229,28 +256,38 @@ class FalAPI_recraft_upscale:
def tensor_to_tempfile(self, tensor):
"""
Convert a ComfyUI IMAGE tensor to a PNG file
Return an open file in "rb" mode (caller must close/delete).
Return the file path (caller must delete).
"""
pil_img = self.tensor_to_pil(tensor)
fd, filename = tempfile.mkstemp(suffix=".png")
os.close(fd)
pil_img.save(filename, format="PNG")
return open(filename, "rb")
return filename
def tensor_to_pil(self, tensor):
"""
Convert tensor to PIL Image: handle (B, H, W, C) or (C, H, W).
"""
print(f"[Recraft Upscale] Converting tensor to PIL - input shape: {tensor.shape}")
if len(tensor.shape) == 4:
tensor = tensor[0] # remove batch dimension
print(f"[Recraft Upscale] Removed batch dimension, new shape: {tensor.shape}")
arr = tensor.cpu().numpy()
print(f"[Recraft Upscale] Converted to numpy array, shape: {arr.shape}")
# If shape is (C, H, W), transpose it
if arr.ndim == 3 and arr.shape[0] <= 4:
arr = np.transpose(arr, (1, 2, 0))
print(f"[Recraft Upscale] Transposed array, new shape: {arr.shape}")
arr = (arr * 255).clip(0, 255).astype("uint8")
return Image.fromarray(arr)
print(f"[Recraft Upscale] Converted to uint8, shape: {arr.shape}")
pil_img = Image.fromarray(arr)
print(f"[Recraft Upscale] Created PIL image: {pil_img.size} mode: {pil_img.mode}")
return pil_img
@classmethod
def IS_CHANGED(cls, **kwargs):