waxwell7/qwen-camera-ctrl

1import argparse2import os3import math4import torch5from PIL import Image6try:7    from diffusers import QwenImageEditPlusPipeline8except ImportError:9    print("Error: QwenImageEditPlusPipeline not found in diffusers.")10    print("Please ensure you have the latest version of diffusers installed.")11    print("pip install git+https://github.com/huggingface/diffusers.git")12    # Fallback for development/testing if not available yet13    QwenImageEditPlusPipeline = None14 15class CameraController:16    def __init__(self, base_model_id="Qwen/Qwen-Image-Edit-2511", lora_path=None, device="cuda"):17        self.device = device18        if device == "cuda" and not torch.cuda.is_available():19            if torch.backends.mps.is_available():20                self.device = "mps"21            else:22                self.device = "cpu"23        24        print(f"Initializing pipeline on {self.device}...")25        26        if QwenImageEditPlusPipeline is None:27            raise ImportError("QwenImageEditPlusPipeline is not available.")28 29        dtype = torch.bfloat16 if self.device != "cpu" else torch.float3230        31        self.pipeline = QwenImageEditPlusPipeline.from_pretrained(32            base_model_id,33            torch_dtype=dtype34        )35        36        if lora_path:37            print(f"Loading LoRA from {lora_path}...")38            self.pipeline.load_lora_weights(lora_path)39            40        self.pipeline.to(self.device)41        print("Pipeline ready.")42 43    def _get_azimuth_desc(self, angle):44        # Normalize angle to 0-36045        angle = angle % 36046        47        # Map to closest 45-degree increment48        # 0, 45, 90, 135, 180, 225, 270, 31549        azimuths = {50            0: "front view",51            45: "front-right quarter view",52            90: "right side view",53            135: "back-right quarter view",54            180: "back view",55            225: "back-left quarter view",56            270: "left side view",57            315: "front-left quarter view"58        }59        60        closest_angle = min(azimuths.keys(), key=lambda x: min(abs(x - angle), abs(x - angle + 360), abs(x - angle - 360)))61        return azimuths[closest_angle]62 63    def _get_elevation_desc(self, angle):64        # Map to closest supported elevation: -30, 0, 30, 6065        elevations = {66            -30: "low-angle shot",67            0: "eye-level shot",68            30: "elevated shot",69            60: "high-angle shot"70        }71        72        closest_angle = min(elevations.keys(), key=lambda x: abs(x - angle))73        return elevations[closest_angle]74 75    def _get_distance_desc(self, distance):76        if isinstance(distance, str):77            distance = distance.lower()78            if "close" in distance: return "close-up"79            if "medium" in distance: return "medium shot"80            if "wide" in distance: return "wide shot"81            return "medium shot" # Default82            83        # If numeric factor84        # 0.6 -> close-up, 1.0 -> medium, 1.8 -> wide85        distances = {86            0.6: "close-up",87            1.0: "medium shot",88            1.8: "wide shot"89        }90        closest_dist = min(distances.keys(), key=lambda x: abs(x - distance))91        return distances[closest_dist]92 93    def generate(self, image_path, azimuth, elevation, distance, output_path="output.png", seed=None):94        if not os.path.exists(image_path):95            raise FileNotFoundError(f"Image not found: {image_path}")96            97        image = Image.open(image_path).convert("RGB")98        99        azimuth_desc = self._get_azimuth_desc(azimuth)100        elevation_desc = self._get_elevation_desc(elevation)101        distance_desc = self._get_distance_desc(distance)102        103        prompt = f"<sks> {azimuth_desc} {elevation_desc} {distance_desc}"104        print(f"Generated prompt: {prompt}")105        106        generator = None107        if seed is not None:108            generator = torch.manual_seed(seed)109            110        inputs = {111            "image": [image], # The pipeline expects a list of images or single image? Search result used [image1, image2] but README implies single image editing. 112                              # Wait, search result says "image": [image1, image2] for Qwen-Image-Edit-2511? 113                              # Ah, the search result example for 2511 uses TWO images. 114                              # But the README for the LoRA says "Input image matters" (singular). 115                              # The standard Qwen-Image-Edit takes one image. 116                              # Let's check the pipeline signature if possible. 117                              # The search result example 1 uses "image": image (single).118                              # The search result example 3 (2511) uses "image": [image1, image2].119                              # It seems 2511 supports multiple images, but for single image editing, we might just pass one.120                              # I'll pass a single image for now.121            "prompt": prompt,122            "generator": generator,123            "true_cfg_scale": 4.0,124            "negative_prompt": " ",125            "num_inference_steps": 40,126            "guidance_scale": 1.0, 127            "num_images_per_prompt": 1,128        }129        130        # Handling different input types for 'image' based on potential pipeline variations131        # If it fails with list, try single item.132        try:133             output = self.pipeline(**inputs)134        except Exception as e:135            print(f"Error with list input, trying single image: {e}")136            inputs["image"] = image137            output = self.pipeline(**inputs)138 139        output_image = output.images[0]140        output_image.save(output_path)141        print(f"Image saved to {output_path}")142        return output_path143 144def main():145    parser = argparse.ArgumentParser(description="Generate images with specific camera angles using Qwen-Image-Edit-2511")146    parser.add_argument("--image", required=True, help="Path to input image")147    parser.add_argument("--azimuth", type=float, default=0, help="Azimuth angle (0-360, 0 is front)")148    parser.add_argument("--elevation", type=float, default=0, help="Elevation angle (-30 to 60, 0 is eye-level)")149    parser.add_argument("--distance", default="medium", help="Distance (close-up, medium, wide) or factor (0.6, 1.0, 1.8)")150    parser.add_argument("--output", default="output.png", help="Output path")151    parser.add_argument("--lora", default="qwenimageedit/qwen-image-edit-2511-multiple-angles-lora.safetensors", help="Path to LoRA file")152    parser.add_argument("--base_model", default="Qwen/Qwen-Image-Edit-2511", help="Base model ID")153    parser.add_argument("--seed", type=int, default=None, help="Random seed")154    parser.add_argument("--low_vram", action="store_true", help="Enable CPU offloading for low VRAM")155    156    args = parser.parse_args()157    158    try:159        controller = CameraController(base_model_id=args.base_model, lora_path=args.lora)160        161        if args.low_vram:162            print("Enabling CPU offloading for low VRAM...")163            controller.pipeline.enable_model_cpu_offload()164        165        # Parse distance if it's a number166        try:167            dist = float(args.distance)168        except ValueError:169            dist = args.distance170            171        controller.generate(172            args.image, 173            args.azimuth, 174            args.elevation, 175            dist, 176            args.output,177            args.seed178        )179    except Exception as e:180        print(f"Error: {e}")181 182if __name__ == "__main__":183    main()184