Implement ImagetoImage

utils/ai_generator.py

@@ -32,19 +32,26 @@ def generate_ai_image(
     neg_prompt_textbox_value,
     model,
     lora_weights=None,
+    conditioned_image=None,
+    pipeline = "FluxPipeline",
     *args,
     **kwargs
 ):
     seed = random.randint(1, 99999)
     if torch.cuda.is_available():
         print("Local GPU available. Generating image locally.")        
+        if conditioned_image is not None:
+            pipeline = "FluxImg2ImgPipeline"
         return generate_ai_image_local(
             map_option,
             prompt_textbox_value,
             neg_prompt_textbox_value,
             model,
             lora_weights=lora_weights,
-            seed=seed
+            seed=seed,
+            conditioned_image=conditioned_image,
+            pipeline_name=pipeline,
+            strength=0.5
         )
     else:
         print("No local GPU available. Sending request to Hugging Face API.")

utils/ai_generator_diffusers_flux.py

@@ -1,11 +1,11 @@
 # utils/ai_generator_diffusers_flux.py
 import os
 import torch
+from diffusers import FluxPipeline,FluxImg2ImgPipeline
 import accelerate 
 import transformers
 import safetensors
 import xformers
-from diffusers import FluxPipeline
 from diffusers.utils import load_image
 # from huggingface_hub import hf_hub_download
 from PIL import Image
@@ -31,6 +31,12 @@ warnings.filterwarnings("ignore", message=".*Torch was not compiled with flash a
 #print(torch.__version__)  # Ensure it's 2.0 or newer
 #print(torch.cuda.is_available())  # Ensure CUDA is available
 
+PIPELINE_CLASSES = {
+    "FluxPipeline": FluxPipeline,
+    "FluxImg2ImgPipeline": FluxImg2ImgPipeline
+}
+
+
 def generate_image_from_text(
     text,
     model_name="black-forest-labs/FLUX.1-dev",
@@ -101,15 +107,22 @@ def generate_image_lowmem(
     num_inference_steps=50,
     seed=0,
     true_cfg_scale=1.0,
+    pipeline_name="FluxPipeline",
+    strength=0.75,
     additional_parameters=None
 ):
+    # Retrieve the pipeline class from the mapping
+    pipeline_class = PIPELINE_CLASSES.get(pipeline_name)
+    if not pipeline_class:
+        raise ValueError(f"Unsupported pipeline type '{pipeline_name}'. "
+                        f"Available options: {list(PIPELINE_CLASSES.keys())}")
     device = "cuda" if torch.cuda.is_available() else "cpu"
     print(f"device:{device}\nmodel_name:{model_name}\n")
     print(f"\n {get_torch_info()}\n")
     # Disable gradient calculations
     with torch.no_grad():
         # Initialize the pipeline inside the context manager
-        pipe = FluxPipeline.from_pretrained(
+        pipe = pipeline_class.from_pretrained(
             model_name,
             torch_dtype=torch.bfloat16 if device == "cuda" else torch.bfloat32            
         ).to(device)
@@ -125,7 +138,8 @@ def generate_image_lowmem(
         else:            
             pipe.attn_implementation="flash_attention_2"
             print("\nEnabled flash_attention_2.\n")
-        pipe.enable_vae_tiling()
+        if pipeline_name == "FluxPipeline":
+            pipe.enable_vae_tiling()
         # Load LoRA weights
         if lora_weights:
             for lora_weight in lora_weights:
@@ -163,22 +177,30 @@ def generate_image_lowmem(
         generator = torch.Generator(device=device).manual_seed(seed)
         conditions = []
         if conditioned_image is not None:
-            conditioned_image = crop_and_resize_image(conditioned_image, 1024, 1024)
+            conditioned_image = crop_and_resize_image(conditioned_image, image_width, image_height)
             condition = Condition("subject", conditioned_image)
             conditions.append(condition)
-        if neg_prompt!=None:
-            true_cfg_scale=1.1
+            print(f"\nAdded conditioned image.\n {conditioned_image.size}")
+            # Prepare the parameters for image generation
+            additional_parameters ={
+                "strength": strength,
+                "image": conditioned_image,
+            }
+        else:
+            print("\nNo conditioned image provided.")
+            if neg_prompt!=None:
+                true_cfg_scale=1.1
+            additional_parameters ={
+                "negative_prompt": neg_prompt,
+                "true_cfg_scale": true_cfg_scale,
+            }
         generate_params = {
-            "prompt": text,
-            "negative_prompt": neg_prompt,
-            "true_cfg_scale": true_cfg_scale,
+            "prompt": text,            
             "height": image_height,
             "width": image_width,
             "guidance_scale": guidance_scale,
             "num_inference_steps": num_inference_steps,
-            "generator": generator,
-            "conditions": conditions if conditions else None
-        }
+            "generator": generator,        }
         if additional_parameters:
             generate_params.update(additional_parameters)
         generate_params = {k: v for k, v in generate_params.items() if v is not None}
@@ -206,7 +228,9 @@ def generate_ai_image_local (
     width=896,
     num_inference_steps=50,
     guidance_scale=3.5,
-    seed=777
+    seed=777,
+    pipeline_name="FluxPipeline",
+    strength=0.75,
 ):
     try:
         if map_option != "Prompt":
@@ -246,6 +270,8 @@ def generate_ai_image_local (
         print(f"Guidance Scale: {guidance_scale}")
         print(f"Seed: {seed}")
         print(f"Additional Parameters: {additional_parameters}")
+        print(f"Conditioned Image: {conditioned_image}")
+        print(f"pipeline: {pipeline_name}")
         image = generate_image_lowmem(
             text=prompt,
             model_name=model,
@@ -257,6 +283,7 @@ def generate_ai_image_local (
             guidance_scale=guidance_scale,
             num_inference_steps=num_inference_steps,
             seed=seed,
+            pipeline_name=pipeline_name,
             additional_parameters=additional_parameters
         )
         with NamedTemporaryFile(delete=False, suffix=".png") as tmp: