Test image2image

app.py

@@ -3,8 +3,7 @@ import gradio as gr
 import numpy as np
 import random
 import os
-# from diffusers import AutoPipelineForText2Image
-from diffusers import FluxPipeline
+from diffusers import FluxImg2ImgPipeline
 import torch
 from huggingface_hub import snapshot_download
 
@@ -13,7 +12,7 @@ SPACE_PASSWORD = 'KR_4dmin'
 
 snapshot_download(repo_id="Roomie/xavyy", cache_dir='./')
 
-pipeline = FluxPipeline.from_pretrained(
+pipeline = FluxImg2ImgPipeline.from_pretrained(
     'black-forest-labs/FLUX.1-schnell', torch_dtype=torch.bfloat16).to('cuda')
 pipeline.load_lora_weights('Roomie/xavyy', weight_name='xavyy.safetensors')
 
@@ -33,7 +32,7 @@ def infer(prompt, refer_image, height, width):
         # num_inference_steps=num_inference_steps,
         width=width,
         height=height,
-        init_image=refer_image
+        image=refer_image
         # generator=generator
     ).images[0]