shaping latents

app.py

@@ -8,37 +8,18 @@ from torchvision import transforms
 from diffusers import DiffusionPipeline
 
 # Define constants
+dtype = torch.bfloat16
+device = "cuda" if torch.cuda.is_available() else "cpu"
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 2048
-MIN_IMAGE_SIZE = 256
-DEFAULT_IMAGE_SIZE = 1024
-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"
-if device == "cpu":
-    print("Warning: Running on CPU. This may be very slow.")
-
-dtype = torch.float16 if device == "cuda" else torch.float32
-
-def load_model():
-    try:
-        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()
+pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=dtype).to(device)
 
-def preprocess_image(image, target_size):
+def preprocess_image(image):
     # Preprocess the image for the VAE
     preprocess = transforms.Compose([
-        transforms.Resize(target_size, interpolation=transforms.InterpolationMode.LANCZOS),
+        transforms.Resize((512, 512)),  # Adjust the size as needed
         transforms.ToTensor(),
         transforms.Normalize([0.5], [0.5])
     ])
@@ -51,76 +32,51 @@ def encode_image(image, vae):
         latents = vae.encode(image).latent_dist.sample() * 0.18215
     return latents
 
-def validate_inputs(prompt, width, height, num_inference_steps):
-    if not prompt or len(prompt) > MAX_PROMPT_LENGTH:
-        raise ValueError(f"Prompt must be between 1 and {MAX_PROMPT_LENGTH} characters.")
-    if width % 8 != 0 or height % 8 != 0:
-        raise ValueError("Width and height must be divisible by 8.")
-    if width < MIN_IMAGE_SIZE or width > MAX_IMAGE_SIZE or height < MIN_IMAGE_SIZE or height > MAX_IMAGE_SIZE:
-        raise ValueError(f"Image dimensions must be between {MIN_IMAGE_SIZE} and {MAX_IMAGE_SIZE}.")
-    if num_inference_steps < 1 or num_inference_steps > 50:
-        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, strength=0.8, progress=gr.Progress(track_tqdm=True)):
-    try:
-        validate_inputs(prompt, width, height, num_inference_steps)
-        
-        if randomize_seed:
-            seed = random.randint(0, MAX_SEED)
-        generator = torch.Generator(device=device).manual_seed(seed)
-        
-        # Ensure max_sequence_length is not more than 256
-        max_sequence_length = min(MAX_PROMPT_LENGTH, len(prompt))
+def infer(prompt, init_image=None, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=4, progress=gr.Progress(track_tqdm=True)):
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    generator = torch.Generator().manual_seed(seed)
+    
+    if init_image is not None:
+        # Process img2img
+        init_image = init_image.convert("RGB")
+        init_image = preprocess_image(init_image)
+        latents = encode_image(init_image, pipe.vae)
+        # Ensure latents are correctly shaped and adjusted
+        latents = torch.nn.functional.interpolate(latents, size=(height // 8, width // 8))
+        latents = latents * 0.18215  # Adjust latent scaling factor if necessary
         
-        if init_image is not None:
-            # Process img2img
-            init_image = init_image.convert("RGB")
-            init_image = preprocess_image(init_image, (height, width))
-            
-            # Encode the image using the VAE
-            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,
-                width=width,
-                num_inference_steps=num_inference_steps,
-                generator=generator,
-                guidance_scale=0.0,
-                latents=latents,
-                max_sequence_length=max_sequence_length
-            ).images[0]
-        else:
-            # Process text2img
-            image = pipe(
-                prompt=prompt,
-                height=height,
-                width=width,
-                num_inference_steps=num_inference_steps,
-                generator=generator,
-                guidance_scale=0.0,
-                max_sequence_length=max_sequence_length
-            ).images[0]
+        # Ensure latents are reshaped to match the expected input dimensions of the model
+        latents = latents.view(1, -1, height // 8, width // 8)
         
-        return image, seed
-    except Exception as e:
-        raise gr.Error(str(e))
+        image = pipe(
+            prompt=prompt,
+            height=height,
+            width=width,
+            num_inference_steps=num_inference_steps,
+            generator=generator,
+            guidance_scale=0.0,
+            latents=latents
+        ).images[0]
+    else:
+        # Process text2img
+        image = pipe(
+            prompt=prompt,
+            height=height,
+            width=width,
+            num_inference_steps=num_inference_steps,
+            generator=generator,
+            guidance_scale=0.0
+        ).images[0]
+    
+    return image, seed
 
 # Define example prompts
 examples = [
     "a tiny astronaut hatching from an egg on the moon",
     "a cat holding a sign that says hello world",
     "an anime illustration of a wiener schnitzel",
-    "a surreal landscape with floating islands and waterfalls",
-    "a steampunk-inspired cityscape at sunset"
 ]
 
 # CSS styling for the Japanese-inspired interface
@@ -173,7 +129,7 @@ with gr.Blocks(css=css) as demo:
                 label="Prompt",
                 show_label=False,
                 max_lines=1,
-                placeholder=f"Enter your prompt (max {MAX_PROMPT_LENGTH} characters)",
+                placeholder="Enter your prompt",
                 container=False,
             )
             run_button = gr.Button("Run", scale=0)
@@ -195,17 +151,17 @@ with gr.Blocks(css=css) as demo:
             with gr.Row():
                 width = gr.Slider(
                     label="Width",
-                    minimum=MIN_IMAGE_SIZE,
+                    minimum=256,
                     maximum=MAX_IMAGE_SIZE,
-                    step=8,
-                    value=DEFAULT_IMAGE_SIZE,
+                    step=32,
+                    value=1024,
                 )
                 height = gr.Slider(
                     label="Height",
-                    minimum=MIN_IMAGE_SIZE,
+                    minimum=256,
                     maximum=MAX_IMAGE_SIZE,
-                    step=8,
-                    value=DEFAULT_IMAGE_SIZE,
+                    step=32,
+                    value=1024,
                 )
 
             with gr.Row():
@@ -216,13 +172,6 @@ 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,
@@ -235,9 +184,8 @@ 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, strength],
+        inputs=[prompt, init_image, seed, randomize_seed, width, height, num_inference_steps],
         outputs=[result, seed]
     )
 
-if __name__ == "__main__":
-    demo.launch()
+demo.launch()