Update app.py

app.py

@@ -5,11 +5,22 @@ import spaces
 import torch
 from diffusers import DiffusionPipeline, FlowMatchEulerDiscreteScheduler, FluxTransformer2DModel, FluxPipeline
 from transformers import CLIPTextModel, CLIPTokenizer,T5EncoderModel, T5TokenizerFast
+from huggingface_hub import hf_hub_download
 
 dtype = torch.bfloat16
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-pipe = FluxPipeline.from_pretrained("sayakpaul/FLUX.1-merged", torch_dtype=torch.bfloat16).to(device)
+repo_name = "ByteDance/Hyper-SD"
+ckpt_name = "Hyper-FLUX.1-dev-8steps-lora.safetensors"
+hyper_lora = hf_hub_download(repo_name, ckpt_name)
+
+pipe = FluxPipeline.from_pretrained(base_model_id, token="xxx")
+pipe.load_lora_weights(hf_hub_download(repo_name, ckpt_name))
+pipe.fuse_lora(lora_scale=0.125)
+pipe.to("cuda", dtype=torch.float16)
+
+
+# pipe = FluxPipeline.from_pretrained("sayakpaul/FLUX.1-merged", torch_dtype=torch.bfloat16).to(device)
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 2048
@@ -20,14 +31,24 @@ def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, guidan
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     generator = torch.Generator().manual_seed(seed)
-    image = pipe(
-        prompt = prompt, 
-        width = width,
-        height = height,
-        num_inference_steps = num_inference_steps, 
-        generator = generator,
-        guidance_scale=guidance_scale
-    ).images[0] 
+    # image = pipe(
+    #     prompt = prompt, 
+    #     width = width,
+    #     height = height,
+    #     num_inference_steps = num_inference_steps, 
+    #     generator = generator,
+    #     guidance_scale=guidance_scale
+    # ).images[0] 
+
+    image = pipe(prompt=prompt, 
+                num_inference_steps=num_inference_steps,
+                guidance_scale=guidance_scale,
+                height=height,
+                width=width,
+                max_sequence_length=256,
+                num_inference_steps=8, 
+                generator = generator,
+                guidance_scale=guidance_scale).images[0]
     return image, seed
  
 examples = [