feat: resize b4 inf

app_v4.py

@@ -56,6 +56,13 @@ try:
 except Exception as e:
     print(f"Failed to dump env info: {e}")
 
+def resize_image_to_max_side(image: Image, max_side_length=1024) -> Image:
+    width, height = image.size
+    ratio = min(max_side_length / width, max_side_length / height)
+    new_size = (int(width * ratio), int(height * ratio))
+    resized_image = image.resize(new_size, Image.BILINEAR)
+    return resized_image
+
 @spaces.GPU(duration=6, progress=gr.Progress(track_tqdm=True))
 @torch.no_grad()
 def generate_image(prompt, scale, steps, control_image, controlnet_conditioning_scale, guidance_scale, seed, guidance_end):
@@ -94,12 +101,14 @@ def combine_caption_focus(caption, focus):
     except Exception as e:
         print(f"Error combining caption and focus: {e}")
         return "highly detailed photo, raw photography."
-
 def generate_caption(control_image):
     try:
         if control_image is None:
             return "Waiting for control image..."
         
+        # Resize the image to a maximum longest side of 1024 pixels
+        control_image = resize_image_to_max_side(control_image, max_side_length=1024)
+
         # Generate a detailed caption
         mcaption = model.caption(control_image, length="short")
         detailed_caption = mcaption["caption"]
@@ -116,56 +125,50 @@ def generate_focus(control_image, focus_list):
             return None
         if focus_list is None:
             return ""
+        
+        # Resize the image to a maximum longest side of 1024 pixels
+        control_image = resize_image_to_max_side(control_image, max_side_length=1024)
+
         # Generate a detailed caption
         focus_query = model.query(control_image, "Please provide a concise but illustrative description of the following area(s) of focus: " + focus_list)
         focus_description = focus_query["answer"]
         print(f"Areas of focus: {focus_description}") 
-
         return focus_description
     except Exception as e:
         print(f"Error generating focus: {e}")
         return "highly detailed photo, raw photography."
 
-def process_image(control_image, user_prompt, system_prompt, scale, steps, 
-                controlnet_conditioning_scale, guidance_scale, seed, 
-                guidance_end, temperature, top_p, max_new_tokens, log_prompt):
-    # Initialize with empty caption
-    final_prompt = user_prompt.strip()
-    # If no user prompt provided, generate a caption first
-    if not final_prompt:
-        # Generate a detailed caption
-        print("Generating caption...")
-        mcaption = model.caption(control_image, length="normal")
-        detailed_caption = mcaption["caption"]
-        final_prompt = detailed_caption
-        yield f"Using caption: {final_prompt}", None, final_prompt
-    
-    # Show the final prompt being used
-    yield f"Generating with: {final_prompt}", None, final_prompt
-    
-    # Generate the image
-    try:
-        image = generate_image(
-            prompt=final_prompt,
-            scale=scale,
-            steps=steps,
+@spaces.GPU(duration=6, progress=gr.Progress(track_tqdm=True))
+@torch.no_grad()
+def generate_image(prompt, scale, steps, control_image, controlnet_conditioning_scale, guidance_scale, seed, guidance_end):
+    generator = torch.Generator().manual_seed(seed)
+    # Load control image
+    control_image = load_image(control_image)
+
+    # Resize the image to a maximum longest side of 1024 pixels
+    control_image = resize_image_to_max_side(control_image, max_side_length=1024)
+
+    w, h = control_image.size
+    w = w - w % 32
+    h = h - h % 32
+    control_image = control_image.resize((int(w * scale), int(h * scale)), resample=2)  # Resample.BILINEAR
+    print("Size to: " + str(control_image.size[0]) + ", " + str(control_image.size[1]))
+    print(f"PromptLog: {repr(prompt)}")
+    with torch.inference_mode():
+        image = pipe(
+            generator=generator,
+            prompt=prompt,
             control_image=control_image,
             controlnet_conditioning_scale=controlnet_conditioning_scale,
+            num_inference_steps=steps,
             guidance_scale=guidance_scale,
-            seed=seed,
-            guidance_end=guidance_end
-        )
-        
-        try:
-            debug_img = Image.open(image.save("/tmp/" + str(seed) + "output.png"))
-            save_image("/tmp/" + str(seed) + "output.png", debug_img)
-        except Exception as e:
-            print("Error 160: " + str(e))
-        log_params(final_prompt, scale, steps, controlnet_conditioning_scale, guidance_scale, seed, guidance_end, control_image, image)
-        yield f"Completed! Used prompt: {final_prompt}", image, final_prompt
-    except Exception as e:
-        print("Error: " + str(e))
-        yield f"Error: {str(e)}", None, None
+            height=control_image.size[1],
+            width=control_image.size[0],
+            control_guidance_start=0.0,
+            control_guidance_end=guidance_end,
+        ).images[0]
+        # print("Type: " + str(type(image)))
+    return image
     
 with gr.Blocks(title="FLUX Turbo Upscaler", fill_height=True) as demo:
     gr.Markdown("⚠️ WIP SPACE - UNFINISHED & BUGGY")