Views
No views yet

1import torch
2from PIL import Image
3from diffusers import DiffusionPipeline
4import requests
5from io import BytesIO
6
7# Configuration
8MODEL_NAME = "/hpc2hdd/home/sfei285/Project/flymyai-lora-trainer/Qwen-Image-Edit" # Replace with your model
9PROMPT = "Use the nano-banana model to create a 1/7 scale commercialized figure of thecharacter in the illustration, in a realistic style and environment. Place the figure on a computer desk, using a circular transparent acrylic base without any text.On the computer screen, display the ZBrush modeling process of the figure. Next to the computer screen, place a BANDAI-style toy packaging box printed with the original artwork."
10NEGATIVE_PROMPT = "lowres, bad anatomy, error body, error arm, error hand, error fingers, error leg, error foot, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, out of frame, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry"
11
12# Device setup
13device = "cuda"
14torch_dtype = torch.bfloat16
15
16print(f"Using device: {device}")
17
18url = "https://huggingface.co/W2GenAI/Nano-Banana-Figure-LoRA/resolve/main/example.jpg"
19response = requests.get(url, timeout=30)
20input_image = Image.open(BytesIO(response.content)).convert("RGB")
21input_image.save("input.jpg")
22
23# Load model
24print(f"Loading model: {MODEL_NAME}")
25pipe = DiffusionPipeline.from_pretrained(MODEL_NAME, torch_dtype=torch_dtype)
26
27# Load LoRA weights
28LORA_WEIGHTS = "W2GenAI/Nano-Banana-Figure-LoRA"
29print(f"Loading LoRA weights: {LORA_WEIGHTS}")
30pipe.load_lora_weights(LORA_WEIGHTS, adapter_name="lora")
31
32# Move to device
33pipe.to(device, dtype=torch_dtype)
34pipe.enable_model_cpu_offload()
35print("Model ready for inference")
36
37# Generate
38generator = torch.Generator(device=device).manual_seed(0)
39with torch.inference_mode():
40 result = pipe(
41 image=input_image,
42 prompt=PROMPT,
43 negative_prompt=NEGATIVE_PROMPT,
44 num_inference_steps=50,
45 true_cfg_scale=4.0,
46 generator=generator,
47 )
48
49# Save result
50result.images[0].save("output.jpg")
51print("Result saved to output.jpg")