Upload app.py

app.py

@@ -24,9 +24,11 @@ from diffusers import AutoencoderKL
 from torch import Tensor, nn
 from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5Tokenizer
 from safetensors.torch import load_file
+from diffusers import FluxTransformer2DModel
+from pipeline_flux_de_distill import FluxPipeline
 dtype = torch.bfloat16
-from huggingface_hub import snapshot_download
-model_path  = snapshot_download(repo_id="nyanko7/flux-dev-de-distill")
+#from huggingface_hub import snapshot_download
+#model_path  = snapshot_download(repo_id="nyanko7/flux-dev-de-distill")
 device = "cuda" if torch.cuda.is_available() else "cpu"
 # ---------------- Encoders ----------------
 
@@ -65,10 +67,14 @@ class HFEmbedder(nn.Module):
         return outputs[self.output_key]
     
 
+model_path = "camenduru/FLUX.1-dev-diffusers"
+transformer = FluxTransformer2DModel.from_pretrained("InstantX/flux-dev-de-distill-diffusers", torch_dtype=torch.bfloat16)
+pipeline = FluxPipeline.from_pretrained(model_path, transformer=transformer, torch_dtype=torch.bfloat16).to(device)
+
 #device = "cuda"
-t5 = HFEmbedder("DeepFloyd/t5-v1_1-xxl", max_length=512, torch_dtype=torch.bfloat16).to(device)
-clip = HFEmbedder("openai/clip-vit-large-patch14", max_length=77, torch_dtype=torch.bfloat16).to(device)
-ae = AutoencoderKL.from_pretrained("camenduru/FLUX.1-dev-diffusers", subfolder="vae", torch_dtype=torch.bfloat16).to(device)
+#t5 = HFEmbedder("DeepFloyd/t5-v1_1-xxl", max_length=512, torch_dtype=torch.bfloat16).to(device)
+#clip = HFEmbedder("openai/clip-vit-large-patch14", max_length=77, torch_dtype=torch.bfloat16).to(device)
+#ae = AutoencoderKL.from_pretrained("camenduru/FLUX.1-dev-diffusers", subfolder="vae", torch_dtype=torch.bfloat16).to(device)
 # quantize(t5, weights=qfloat8)
 # freeze(t5)
 
@@ -655,9 +661,10 @@ def generate_image(
 ):
     if seed == 0:
         seed = int(random.random() * 1000000)
+    generator = torch.Generator(device=device).manual_seed(seed)
         
     #device = "cuda" if torch.cuda.is_available() else "cpu"
-    torch_device = torch.device(device)
+    """torch_device = torch.device(device)
         
     if do_img2img and init_image is not None:
         init_image = get_image(init_image)
@@ -696,8 +703,18 @@ def generate_image(
 
     x = x.clamp(-1, 1)
     x = rearrange(x[0], "c h w -> h w c")
-    img = Image.fromarray((127.5 * (x + 1.0)).cpu().byte().numpy())
+    img = Image.fromarray((127.5 * (x + 1.0)).cpu().byte().numpy())"""
     
+    img = pipeline(
+        prompt=prompt,
+        negative_prompt=neg_prompt,
+        guidance_scale=guidance,
+        num_inference_steps=num_steps,
+        width=width,
+        height=height,
+        generator=generator,
+    ).images[0]
+
     return img, seed
 
 def create_demo():