Update app.py

app.py

@@ -12,42 +12,20 @@ from src.pipeline import FluxPipeline
 from src.transformer_flux import FluxTransformer2DModel
 from src.lora_helper import set_single_lora, set_multi_lora, unset_lora
 
-
-class ImageProcessor:
-    def __init__(self, path):
-        device = "cuda"
-        self.pipe = FluxPipeline.from_pretrained(path, torch_dtype=torch.bfloat16, device=device)
-        transformer = FluxTransformer2DModel.from_pretrained(path, subfolder="transformer", torch_dtype=torch.bfloat16, device=device)
-        self.pipe.transformer = transformer
-        self.pipe.to(device)
-        
-    def clear_cache(self, transformer):
-        for name, attn_processor in transformer.attn_processors.items():
-            attn_processor.bank_kv.clear()
-            
-    def process_image(self, prompt='', subject_imgs=[], spatial_imgs=[], height=768, width=768, output_path=None, seed=42):
-        image = self.pipe(
-            prompt,
-            height=int(height),
-            width=int(width),
-            guidance_scale=3.5,
-            num_inference_steps=25,
-            max_sequence_length=512,
-            generator=torch.Generator("cpu").manual_seed(seed), 
-            subject_images=subject_imgs,
-            spatial_images=spatial_imgs,
-            cond_size=512,
-        ).images[0]
-        self.clear_cache(self.pipe.transformer)
-        if output_path:
-            image.save(output_path)
-        return image
-
 # Initialize the image processor
 base_path = "black-forest-labs/FLUX.1-dev"    
 lora_base_path = "./models"
 style_lora_base_path = "Shakker-Labs"
-processor = ImageProcessor(base_path)
+
+
+pipe = FluxPipeline.from_pretrained(base_path, torch_dtype=torch.bfloat16)
+transformer = FluxTransformer2DModel.from_pretrained(base_path, subfolder="transformer", torch_dtype=torch.bfloat16)
+pipe.transformer = transformer
+pipe.to(device)
+
+def clear_cache(transformer):
+    for name, attn_processor in transformer.attn_processors.items():
+        attn_processor.bank_kv.clear()
 
 # Define the Gradio interface
 @spaces.GPU()
@@ -59,29 +37,41 @@ def single_condition_generate_image(prompt, subject_img, spatial_img, height, wi
         lora_path = os.path.join(lora_base_path, "pose.safetensors")
     elif control_type == "inpainting":
         lora_path = os.path.join(lora_base_path, "inpainting.safetensors")
-    set_single_lora(processor.pipe.transformer, lora_path, lora_weights=[1], cond_size=512)
+    set_single_lora(pipe.transformer, lora_path, lora_weights=[1], cond_size=512)
     
     # Set the style LoRA
     if style_lora=="None":
         pass
     else:
         if style_lora == "Simple_Sketch":
-            processor.pipe.unload_lora_weights()
+            pipe.unload_lora_weights()
             style_lora_path = os.path.join(style_lora_base_path, "FLUX.1-dev-LoRA-Children-Simple-Sketch")
-            processor.pipe.load_lora_weights(style_lora_path, weight_name="FLUX-dev-lora-children-simple-sketch.safetensors")
+            pipe.load_lora_weights(style_lora_path, weight_name="FLUX-dev-lora-children-simple-sketch.safetensors")
         if style_lora == "Text_Poster":
-            processor.pipe.unload_lora_weights()
+            pipe.unload_lora_weights()
             style_lora_path = os.path.join(style_lora_base_path, "FLUX.1-dev-LoRA-Text-Poster")
-            processor.pipe.load_lora_weights(style_lora_path, weight_name="FLUX-dev-lora-Text-Poster.safetensors")
+            pipe.load_lora_weights(style_lora_path, weight_name="FLUX-dev-lora-Text-Poster.safetensors")
         if style_lora == "Vector_Style":
-            processor.pipe.unload_lora_weights()
+            pipe.unload_lora_weights()
             style_lora_path = os.path.join(style_lora_base_path, "FLUX.1-dev-LoRA-Vector-Journey")
-            processor.pipe.load_lora_weights(style_lora_path, weight_name="FLUX-dev-lora-Vector-Journey.safetensors")
+            pipe.load_lora_weights(style_lora_path, weight_name="FLUX-dev-lora-Vector-Journey.safetensors")
 
     # Process the image
     subject_imgs = [subject_img] if subject_img else []
     spatial_imgs = [spatial_img] if spatial_img else []
-    image = processor.process_image(prompt=prompt, subject_imgs=subject_imgs, spatial_imgs=spatial_imgs, height=height, width=width, seed=seed)
+    image = pipe(
+        prompt,
+        height=int(height),
+        width=int(width),
+        guidance_scale=3.5,
+        num_inference_steps=25,
+        max_sequence_length=512,
+        generator=torch.Generator("cpu").manual_seed(seed), 
+        subject_images=subject_imgs,
+        spatial_images=spatial_imgs,
+        cond_size=512,
+    ).images[0]
+    clear_cache(pipe.transformer)
     return image
 
 # Define the Gradio interface
@@ -89,12 +79,24 @@ def single_condition_generate_image(prompt, subject_img, spatial_img, height, wi
 def multi_condition_generate_image(prompt, subject_img, spatial_img, height, width, seed):
     subject_path = os.path.join(lora_base_path, "subject.safetensors")
     inpainting_path = os.path.join(lora_base_path, "inpainting.safetensors")
-    set_multi_lora(processor.pipe.transformer, [subject_path, inpainting_path], lora_weights=[[1],[1]],cond_size=512)
+    set_multi_lora(pipe.transformer, [subject_path, inpainting_path], lora_weights=[[1],[1]],cond_size=512)
 
     # Process the image
     subject_imgs = [subject_img] if subject_img else []
     spatial_imgs = [spatial_img] if spatial_img else []
-    image = processor.process_image(prompt=prompt, subject_imgs=subject_imgs, spatial_imgs=spatial_imgs, height=height, width=width, seed=seed)
+    image = pipe(
+        prompt,
+        height=int(height),
+        width=int(width),
+        guidance_scale=3.5,
+        num_inference_steps=25,
+        max_sequence_length=512,
+        generator=torch.Generator("cpu").manual_seed(seed), 
+        subject_images=subject_imgs,
+        spatial_images=spatial_imgs,
+        cond_size=512,
+    ).images[0]
+    clear_cache(pipe.transformer)
     return image
 
 # Define the Gradio interface components