Upload 2 files

handler.py

@@ -2,7 +2,6 @@ import os
 from typing import Any, Dict
 
 from diffusers import FluxPipeline, FluxTransformer2DModel, AutoencoderKL, TorchAoConfig
-from diffusers.utils.remote_utils import remote_decode
 from PIL import Image
 import torch
 
@@ -27,10 +26,9 @@ class EndpointHandler:
         #repo_id = "NoMoreCopyright/FLUX.1-dev-test"
         dtype = torch.bfloat16
         quantization_config = TorchAoConfig("int8wo")
-        #vae = AutoencoderKL.from_pretrained(repo_id, subfolder="vae", torch_dtype=dtype)
+        vae = AutoencoderKL.from_pretrained(repo_id, subfolder="vae", torch_dtype=dtype)
         #transformer = FluxTransformer2DModel.from_pretrained(repo_id, subfolder="transformer", torch_dtype=dtype, quantization_config=quantization_config).to("cuda")
-        #self.pipeline = FluxPipeline.from_pretrained(repo_id, vae=vae, torch_dtype=dtype, quantization_config=quantization_config)
-        self.pipeline = FluxPipeline.from_pretrained(repo_id, vae=None, torch_dtype=dtype, quantization_config=quantization_config)
+        self.pipeline = FluxPipeline.from_pretrained(repo_id, vae=vae, torch_dtype=dtype, quantization_config=quantization_config)
         if IS_COMPILE: self.pipeline = compile_pipeline(self.pipeline)
         self.pipeline.to("cuda")
 
@@ -58,23 +56,13 @@ class EndpointHandler:
         seed = parameters.get("seed", 0)
         generator = torch.manual_seed(seed)
 
-        latent = self.pipeline(  # type: ignore
+        return self.pipeline(  # type: ignore
             prompt,
             height=height,
             width=width,
             guidance_scale=guidance_scale,
             num_inference_steps=num_inference_steps,
             generator=generator,
-            output_type="latent",
-        ).images
-    
-        image = remote_decode(
-            endpoint="https://whhx50ex1aryqvw6.us-east-1.aws.endpoints.huggingface.cloud/",
-            tensor=latent,
-            height=height,
-            width=width,
-            scaling_factor=0.3611,
-            shift_factor=0.1159,
-        )
+            output_type="pil",
+        ).images[0]
 
-        return image

requirements.txt

@@ -2,7 +2,7 @@ huggingface_hub
 torch
 torchvision
 torchao
-git+https://github.com/huggingface/diffusers
+diffusers
 peft
 accelerate
 transformers
@@ -11,4 +11,5 @@ scipy
 Pillow
 sentencepiece
 protobuf
-pytorch-lightning
+pytorch-lightning
+xformers