Spaces:
Runtime error
Runtime error
| import random | |
| import gradio as gr | |
| import spaces | |
| import torch | |
| from PIL import Image | |
| from diffusers import ZImageImg2ImgPipeline | |
| # ============================================================ | |
| # CONFIG | |
| # ============================================================ | |
| MODEL_ID = "Tongyi-MAI/Z-Image-Turbo" | |
| # IMPORTANT: | |
| # Z-Image strength=1.0 essentially ignores the input image. | |
| # We intentionally cap the UI at 0.85 for better editing behavior. | |
| MIN_STRENGTH = 0.15 | |
| MAX_STRENGTH = 0.85 | |
| DEFAULT_STRENGTH = 0.55 | |
| NUM_INFERENCE_STEPS = 8 | |
| GUIDANCE_SCALE = 0.0 | |
| MAX_IMAGE_SIZE = 1024 | |
| GPU_DURATION = 180 | |
| # ============================================================ | |
| # MODEL | |
| # ============================================================ | |
| print( | |
| "Loading Z-Image-Turbo...", | |
| flush=True, | |
| ) | |
| pipe = ZImageImg2ImgPipeline.from_pretrained( | |
| MODEL_ID, | |
| torch_dtype=torch.bfloat16, | |
| ) | |
| pipe.to("cuda") | |
| pipe.set_progress_bar_config( | |
| disable=True | |
| ) | |
| print( | |
| "Z-Image-Turbo loaded successfully.", | |
| flush=True, | |
| ) | |
| # ============================================================ | |
| # IMAGE PREPARATION | |
| # ============================================================ | |
| def prepare_image( | |
| image: Image.Image, | |
| ) -> Image.Image: | |
| if image is None: | |
| raise ValueError( | |
| "Please upload an image." | |
| ) | |
| image = image.convert("RGB") | |
| width, height = image.size | |
| if width < 64 or height < 64: | |
| raise ValueError( | |
| "Image must be at least 64×64 pixels." | |
| ) | |
| # Keep the original aspect ratio. | |
| scale = min( | |
| MAX_IMAGE_SIZE / width, | |
| MAX_IMAGE_SIZE / height, | |
| 1.0, | |
| ) | |
| new_width = int(width * scale) | |
| new_height = int(height * scale) | |
| # Z-Image works with dimensions divisible by 8. | |
| new_width = max( | |
| 64, | |
| (new_width // 8) * 8, | |
| ) | |
| new_height = max( | |
| 64, | |
| (new_height // 8) * 8, | |
| ) | |
| if ( | |
| new_width != image.width | |
| or new_height != image.height | |
| ): | |
| image = image.resize( | |
| (new_width, new_height), | |
| Image.Resampling.LANCZOS, | |
| ) | |
| return image | |
| # ============================================================ | |
| # PROMPT HANDLING | |
| # ============================================================ | |
| def build_edit_prompt( | |
| prompt: str, | |
| strength: float, | |
| ) -> str: | |
| prompt = prompt.strip() | |
| # Lower strength = stronger preservation. | |
| if strength <= 0.35: | |
| prefix = ( | |
| "Edit the provided image while preserving " | |
| "the original subject, identity, composition, " | |
| "camera angle, pose, framing, and overall " | |
| "structure. Make only the requested visual " | |
| "changes. " | |
| ) | |
| elif strength <= 0.60: | |
| prefix = ( | |
| "Transform the provided image according to " | |
| "the request while preserving the main subject, " | |
| "recognizable identity, composition, and " | |
| "important visual structure. " | |
| ) | |
| else: | |
| prefix = ( | |
| "Creatively transform the provided image " | |
| "according to the request while keeping the " | |
| "main subject and important visual identity " | |
| "recognizable. Do not completely replace the " | |
| "original concept unless explicitly requested. " | |
| ) | |
| return prefix + prompt | |
| # ============================================================ | |
| # GENERATION | |
| # ============================================================ | |
| def generate( | |
| image, | |
| prompt, | |
| strength, | |
| seed, | |
| ): | |
| if image is None: | |
| raise gr.Error( | |
| "Please upload an image first." | |
| ) | |
| if not prompt or not prompt.strip(): | |
| raise gr.Error( | |
| "Please enter a transformation prompt." | |
| ) | |
| # -------------------------------------------------------- | |
| # Strength | |
| # -------------------------------------------------------- | |
| try: | |
| strength = float(strength) | |
| except Exception: | |
| raise gr.Error( | |
| "Invalid transformation strength." | |
| ) | |
| strength = max( | |
| MIN_STRENGTH, | |
| min( | |
| MAX_STRENGTH, | |
| strength, | |
| ), | |
| ) | |
| # -------------------------------------------------------- | |
| # Seed | |
| # -------------------------------------------------------- | |
| try: | |
| seed = int(seed) | |
| except Exception: | |
| seed = -1 | |
| if seed < 0: | |
| seed = random.SystemRandom().randint( | |
| 0, | |
| 2**32 - 1, | |
| ) | |
| # -------------------------------------------------------- | |
| # Generation | |
| # -------------------------------------------------------- | |
| try: | |
| image = prepare_image(image) | |
| final_prompt = build_edit_prompt( | |
| prompt, | |
| strength, | |
| ) | |
| generator = torch.Generator( | |
| device="cuda" | |
| ).manual_seed(seed) | |
| print( | |
| "Generation started | " | |
| f"seed={seed} | " | |
| f"strength={strength}", | |
| flush=True, | |
| ) | |
| print( | |
| f"Prompt: {final_prompt}", | |
| flush=True, | |
| ) | |
| with torch.inference_mode(): | |
| result = pipe( | |
| prompt=final_prompt, | |
| image=image, | |
| strength=strength, | |
| num_inference_steps=NUM_INFERENCE_STEPS, | |
| guidance_scale=GUIDANCE_SCALE, | |
| generator=generator, | |
| ) | |
| output = result.images[0] | |
| print( | |
| "Generation complete | " | |
| f"seed={seed}", | |
| flush=True, | |
| ) | |
| return output, seed | |
| except torch.cuda.OutOfMemoryError: | |
| if torch.cuda.is_available(): | |
| torch.cuda.empty_cache() | |
| raise gr.Error( | |
| "GPU memory was insufficient. " | |
| "Try a smaller input image." | |
| ) | |
| except Exception as exc: | |
| print( | |
| "Generation error:", | |
| repr(exc), | |
| flush=True, | |
| ) | |
| if torch.cuda.is_available(): | |
| torch.cuda.empty_cache() | |
| raise gr.Error( | |
| f"Generation failed: {exc}" | |
| ) | |
| # ============================================================ | |
| # UI | |
| # ============================================================ | |
| CSS = """ | |
| #title { | |
| text-align: center; | |
| } | |
| #subtitle { | |
| text-align: center; | |
| opacity: 0.7; | |
| } | |
| .generate-btn { | |
| min-height: 52px; | |
| font-size: 18px !important; | |
| font-weight: 600 !important; | |
| } | |
| .strength-info { | |
| text-align: center; | |
| opacity: 0.7; | |
| font-size: 13px; | |
| } | |
| """ | |
| with gr.Blocks( | |
| title="IMAGINE — Z-Image Image to Image", | |
| css=CSS, | |
| ) as demo: | |
| gr.Markdown( | |
| "# IMAGINE", | |
| elem_id="title", | |
| ) | |
| gr.Markdown( | |
| "Z-Image-Turbo Image Transformation", | |
| elem_id="subtitle", | |
| ) | |
| with gr.Row(): | |
| # ==================================================== | |
| # INPUT | |
| # ==================================================== | |
| with gr.Column(): | |
| input_image = gr.Image( | |
| type="pil", | |
| label="Upload Image", | |
| ) | |
| prompt = gr.Textbox( | |
| label="Transformation Prompt", | |
| placeholder=( | |
| "Example: Change the clothes to " | |
| "a black formal suit while keeping " | |
| "the person's face and pose." | |
| ), | |
| lines=4, | |
| ) | |
| strength = gr.Slider( | |
| minimum=MIN_STRENGTH, | |
| maximum=MAX_STRENGTH, | |
| value=DEFAULT_STRENGTH, | |
| step=0.05, | |
| label="Transformation Strength", | |
| info=( | |
| "Lower = preserve more • " | |
| "Higher = transform more" | |
| ), | |
| ) | |
| gr.Markdown( | |
| "Preserve ←──────── Balanced ────────→ Transform", | |
| elem_classes="strength-info", | |
| ) | |
| seed = gr.Number( | |
| value=-1, | |
| precision=0, | |
| label="Seed", | |
| info="-1 = random", | |
| ) | |
| generate_button = gr.Button( | |
| "Generate Image", | |
| variant="primary", | |
| elem_classes="generate-btn", | |
| ) | |
| # ==================================================== | |
| # OUTPUT | |
| # ==================================================== | |
| with gr.Column(): | |
| output_image = gr.Image( | |
| type="pil", | |
| label="Generated Image", | |
| ) | |
| output_seed = gr.Number( | |
| label="Used Seed", | |
| interactive=False, | |
| ) | |
| # ======================================================== | |
| # EXAMPLES | |
| # ======================================================== | |
| gr.Examples( | |
| examples=[ | |
| [ | |
| ( | |
| "Change the lighting to warm golden " | |
| "hour lighting while keeping the " | |
| "same person, face, pose, background " | |
| "and composition." | |
| ), | |
| 0.30, | |
| -1, | |
| ], | |
| [ | |
| ( | |
| "Change the person's clothing to a " | |
| "premium black formal outfit while " | |
| "keeping the same face, pose and " | |
| "composition." | |
| ), | |
| 0.50, | |
| -1, | |
| ], | |
| [ | |
| ( | |
| "Transform this scene into a cinematic " | |
| "fantasy environment while keeping " | |
| "the main subject recognizable." | |
| ), | |
| 0.70, | |
| -1, | |
| ], | |
| ], | |
| inputs=[ | |
| prompt, | |
| strength, | |
| seed, | |
| ], | |
| ) | |
| # ======================================================== | |
| # API | |
| # ======================================================== | |
| generate_button.click( | |
| fn=generate, | |
| inputs=[ | |
| input_image, | |
| prompt, | |
| strength, | |
| seed, | |
| ], | |
| outputs=[ | |
| output_image, | |
| output_seed, | |
| ], | |
| api_name="generate", | |
| ) | |
| # ============================================================ | |
| # LAUNCH | |
| # ============================================================ | |
| if __name__ == "__main__": | |
| demo.queue( | |
| max_size=20, | |
| default_concurrency_limit=1, | |
| ).launch( | |
| server_name="0.0.0.0", | |
| server_port=7860, | |
| ) |