fix trying to fix image preprocessing

app.py

@@ -8,18 +8,32 @@ 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 = 500
+
+# 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:
+        return DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=dtype).to(device)
+    except Exception as e:
+        raise RuntimeError(f"Failed to load the model: {str(e)}")
 
 # Load the diffusion pipeline
-pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=dtype).to(device)
+pipe = load_model()
 
-def preprocess_image(image):
+def preprocess_image(image, target_size=(512, 512)):
     # Preprocess the image for the VAE
     preprocess = transforms.Compose([
-        transforms.Resize((512, 512)),  # Adjust the size as needed
+        transforms.Resize(target_size, interpolation=transforms.InterpolationMode.LANCZOS),
         transforms.ToTensor(),
         transforms.Normalize([0.5], [0.5])
     ])
@@ -32,44 +46,60 @@ 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=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)
-        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
+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)):
+    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)
+        
+        if init_image is not None:
+            init_image = init_image.convert("RGB")
+            init_image = preprocess_image(init_image, (height, width))
+            latents = encode_image(init_image, pipe.vae)
+            latents = torch.nn.functional.interpolate(latents, size=(height // 8, width // 8), mode='bilinear')
+            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:
+            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
+    except Exception as e:
+        raise gr.Error(str(e))
 
 # 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
@@ -122,7 +152,7 @@ with gr.Blocks(css=css) as demo:
                 label="Prompt",
                 show_label=False,
                 max_lines=1,
-                placeholder="Enter your prompt",
+                placeholder=f"Enter your prompt (max {MAX_PROMPT_LENGTH} characters)",
                 container=False,
             )
             run_button = gr.Button("Run", scale=0)
@@ -144,17 +174,17 @@ with gr.Blocks(css=css) as demo:
             with gr.Row():
                 width = gr.Slider(
                     label="Width",
-                    minimum=256,
+                    minimum=MIN_IMAGE_SIZE,
                     maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024,
+                    step=8,
+                    value=DEFAULT_IMAGE_SIZE,
                 )
                 height = gr.Slider(
                     label="Height",
-                    minimum=256,
+                    minimum=MIN_IMAGE_SIZE,
                     maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024,
+                    step=8,
+                    value=DEFAULT_IMAGE_SIZE,
                 )
 
             with gr.Row():
@@ -181,7 +211,8 @@ with gr.Blocks(css=css) as demo:
         outputs=[result, seed]
     )
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()
 
 
 

requirements.txt

@@ -6,4 +6,5 @@ transformers==4.42.4
 xformers
 sentencepiece
 gradio==4.29.0
-torchvision
+torchvision
+pillow