main push

app.py

@@ -12,7 +12,7 @@ MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 2048
 MIN_IMAGE_SIZE = 256
 DEFAULT_IMAGE_SIZE = 1024
-MAX_PROMPT_LENGTH = 500
+MAX_PROMPT_LENGTH = 256  # Changed to 256 as per FLUX.1-schnell requirements
 
 # Check for GPU availability
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -23,14 +23,19 @@ dtype = torch.float16 if device == "cuda" else torch.float32
 
 def load_model():
     try:
-        return DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=dtype).to(device)
+        pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=dtype)
+        pipe.to(device)
+        pipe.enable_model_cpu_offload()
+        pipe.vae.enable_slicing()
+        pipe.vae.enable_tiling()
+        return pipe
     except Exception as e:
         raise RuntimeError(f"Failed to load the model: {str(e)}")
 
 # Load the diffusion pipeline
 pipe = load_model()
 
-def preprocess_image(image, target_size=(512, 512)):
+def preprocess_image(image, target_size):
     # Preprocess the image for the VAE
     preprocess = transforms.Compose([
         transforms.Resize(target_size, interpolation=transforms.InterpolationMode.LANCZOS),
@@ -57,7 +62,7 @@ def validate_inputs(prompt, width, height, num_inference_steps):
         raise ValueError("Number of inference steps must be between 1 and 50.")
 
 @spaces.GPU()
-def infer(prompt, init_image=None, seed=42, randomize_seed=False, width=DEFAULT_IMAGE_SIZE, height=DEFAULT_IMAGE_SIZE, num_inference_steps=4, progress=gr.Progress(track_tqdm=True)):
+def infer(prompt, init_image=None, seed=42, randomize_seed=False, width=DEFAULT_IMAGE_SIZE, height=DEFAULT_IMAGE_SIZE, num_inference_steps=4, strength=0.8, progress=gr.Progress(track_tqdm=True)):
     try:
         validate_inputs(prompt, width, height, num_inference_steps)
         
@@ -74,13 +79,15 @@ def infer(prompt, init_image=None, seed=42, randomize_seed=False, width=DEFAULT_
             init_image = preprocess_image(init_image, (height, width))
             
             # Encode the image using the VAE
-            with torch.no_grad():
-                init_latents = pipe.vae.encode(init_image).latent_dist.sample(generator=generator)
-                init_latents = 0.18215 * init_latents
+            init_latents = encode_image(init_image, pipe.vae)
             
             # Ensure latents are correctly shaped
             init_latents = torch.nn.functional.interpolate(init_latents, size=(height // 8, width // 8), mode='bilinear', align_corners=False)
             
+            # Add noise to latents
+            noise = torch.randn_like(init_latents)
+            latents = noise + strength * (init_latents - noise)
+            
             image = pipe(
                 prompt=prompt,
                 height=height,
@@ -88,7 +95,7 @@ def infer(prompt, init_image=None, seed=42, randomize_seed=False, width=DEFAULT_
                 num_inference_steps=num_inference_steps,
                 generator=generator,
                 guidance_scale=0.0,
-                latents=init_latents,  # Use latents instead of image
+                latents=latents,
                 max_sequence_length=max_sequence_length
             ).images[0]
         else:
@@ -209,6 +216,13 @@ with gr.Blocks(css=css) as demo:
                     step=1,
                     value=4,
                 )
+                strength = gr.Slider(
+                    label="Strength (for img2img)",
+                    minimum=0.0,
+                    maximum=1.0,
+                    step=0.01,
+                    value=0.8,
+                )
 
         gr.Examples(
             examples=examples,
@@ -221,12 +235,9 @@ with gr.Blocks(css=css) as demo:
     gr.on(
         triggers=[run_button.click, prompt.submit],
         fn=infer,
-        inputs=[prompt, init_image, seed, randomize_seed, width, height, num_inference_steps],
+        inputs=[prompt, init_image, seed, randomize_seed, width, height, num_inference_steps, strength],
         outputs=[result, seed]
     )
 
 if __name__ == "__main__":
-    demo.launch()
-
-
-
+    demo.launch()