from diffusers import StableDiffusionPipeline
import torch

class QwenImageEdit:
    def __init__(self, model_name="Phr00t/Qwen-Image-Edit-Rapid-AIO"):
        self.pipe = StableDiffusionPipeline.from_pretrained(
            model_name,
            torch_dtype=torch.float16,
            revision="fp16",  # if applicable
        ).to("cuda" if torch.cuda.is_available() else "cpu")

    def generate(self, prompt: str, image=None, num_inference_steps=30, guidance_scale=7.5):
        # If image is provided → image-to-image editing, otherwise vanilla text→image
        if image is not None:
            out = self.pipe(
                prompt=prompt,
                init_image=image,
                strength=0.7,
                num_inference_steps=num_inference_steps,
                guidance_scale=guidance_scale
            )
        else:
            out = self.pipe(
                prompt=prompt,
                num_inference_steps=num_inference_steps,
                guidance_scale=guidance_scale
            )
        return out.images[0]