Update app.py

app.py

@@ -6,6 +6,18 @@ import spaces
 from diffusers import DiffusionPipeline
 import torch
 
+
+import numpy as np
+import random
+import spaces
+import torch
+import time
+from diffusers import DiffusionPipeline, AutoencoderTiny
+from diffusers.models.attention_processor import AttnProcessor2_0
+from custom_pipeline import FluxWithCFGPipeline
+
+torch.backends.cuda.matmul.allow_tf32 = True
+
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model_repo_id = "stabilityai/stable-diffusion-3.5-large"
 
@@ -17,6 +29,26 @@ else:
 pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
 pipe = pipe.to(device)
 
+
+
+
+
+dtype = torch.float16
+pipe2 = FluxWithCFGPipeline.from_pretrained(
+    "black-forest-labs/FLUX.1-schnell", torch_dtype=dtype
+)
+pipe2.vae = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype)
+pipe2.to("cuda")
+pipe2.load_lora_weights('hugovntr/flux-schnell-realism', weight_name='schnell-realism_v2.3.safetensors', adapter_name="better")
+pipe2.set_adapters(["better"], adapter_weights=[1.0])
+pipe2.fuse_lora(adapter_name=["better"], lora_scale=1.0)
+pipe2.unload_lora_weights()
+
+torch.cuda.empty_cache()
+
+
+
+
 def adjust_to_nearest_multiple(value, divisor=8):
     """
     Adjusts the input value to the nearest multiple of the divisor.
@@ -68,6 +100,7 @@ def infer(
     num_inference_steps=40,
     progress=gr.Progress(track_tqdm=True),
 ):
+    
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
 
@@ -77,15 +110,24 @@ def infer(
 
     generator = torch.Generator().manual_seed(seed)
 
-    image = pipe(
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        guidance_scale=guidance_scale,
-        num_inference_steps=num_inference_steps,
-        width=width,
-        height=height,
-        generator=generator,
-    ).images[0]
+    if num_inference_steps<=10:
+        image = pipe(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            guidance_scale=guidance_scale,
+            num_inference_steps=num_inference_steps,
+            width=width,
+            height=height,
+            generator=generator,
+        ).images[0]
+    else:
+        img = pipe2.generate_images( 
+            prompt=prompt,
+            width=width,
+            height=height,
+            num_inference_steps=num_inference_steps,
+            generator=generator
+        )
 
     return image, seed