# app_v4.py import gradio as gr import torch import spaces import os import datetime import io import moondream as md from transformers import T5EncoderModel from diffusers import FluxControlNetPipeline from diffusers.utils import load_image from PIL import Image from threading import Thread from typing import Generator from huggingface_hub import CommitScheduler, HfApi, logging from debug import log_params, scheduler, save_image logging.set_verbosity_debug() from model_loader import safe_model_load # Ensure device is set DEVICE = "cuda" if torch.cuda.is_available() else "cpu" MAX_SEED = 1000000 huggingface_token = os.getenv("HUGGINFACE_TOKEN") md_api_key = os.getenv("MD_KEY") model = md.vl(api_key=md_api_key) try: # Set max memory usage for ZeroGPU torch.cuda.set_per_process_memory_fraction(1.0) torch.set_float32_matmul_precision("high") except Exception as e: print(f"Error setting memory usage: {e}") text_encoder_2_unquant = T5EncoderModel.from_pretrained( "LPX55/FLUX.1-merged_uncensored", subfolder="text_encoder_2", torch_dtype=torch.bfloat16, token=huggingface_token ) pipe = FluxControlNetPipeline.from_pretrained( "LPX55/FLUX.1M-8step_upscaler-cnet", torch_dtype=torch.bfloat16, text_encoder_2=text_encoder_2_unquant, token=huggingface_token ) pipe.to("cuda") @spaces.GPU(duration=12) @torch.no_grad() def generate_image(prompt, scale, steps, control_image, controlnet_conditioning_scale, guidance_scale, seed, guidance_end): generator = torch.Generator().manual_seed(seed) # Load control image control_image = load_image(control_image) w, h = control_image.size w = w - w % 32 h = h - h % 32 control_image = control_image.resize((int(w * scale), int(h * scale)), resample=2) # Resample.BILINEAR print("Size to: " + str(control_image.size[0]) + ", " + str(control_image.size[1])) print(f"PromptLog: {repr(prompt)}") with torch.inference_mode(): image = pipe( generator=generator, prompt=prompt, control_image=control_image, controlnet_conditioning_scale=controlnet_conditioning_scale, num_inference_steps=steps, guidance_scale=guidance_scale, height=control_image.size[1], width=control_image.size[0], control_guidance_start=0.0, control_guidance_end=guidance_end, ).images[0] # print("Type: " + str(type(image))) return image def combine_caption_focus(caption, focus): try: if caption is None: caption = "" if focus is None: focus = "highly detailed photo, raw photography." return (str(caption) + "\n\n" + str(focus)).strip() except Exception as e: print(f"Error combining caption and focus: {e}") return "highly detailed photo, raw photography." def generate_caption(control_image): try: if control_image is None: return "Waiting for control image..." # Generate a detailed caption mcaption = model.caption(control_image, length="short") detailed_caption = mcaption["caption"] print(f"Detailed caption: {detailed_caption}") return detailed_caption except Exception as e: print(f"Error generating caption: {e}") return "A detailed photograph" def generate_focus(control_image, focus_list): try: if control_image is None: return None if focus_list is None: return "" # Generate a detailed caption focus_query = model.query(control_image, "Please provide a concise but illustrative description of the following area(s) of focus: " + focus_list) focus_description = focus_query["answer"] print(f"Areas of focus: {focus_description}") return focus_description except Exception as e: print(f"Error generating focus: {e}") return "highly detailed photo, raw photography." def process_image(control_image, user_prompt, system_prompt, scale, steps, controlnet_conditioning_scale, guidance_scale, seed, guidance_end, temperature, top_p, max_new_tokens, log_prompt): # Initialize with empty caption final_prompt = user_prompt.strip() # If no user prompt provided, generate a caption first if not final_prompt: # Generate a detailed caption print("Generating caption...") mcaption = model.caption(control_image, length="normal") detailed_caption = mcaption["caption"] final_prompt = detailed_caption yield f"Using caption: {final_prompt}", None, final_prompt # Show the final prompt being used yield f"Generating with: {final_prompt}", None, final_prompt # Generate the image try: image = generate_image( prompt=final_prompt, scale=scale, steps=steps, control_image=control_image, controlnet_conditioning_scale=controlnet_conditioning_scale, guidance_scale=guidance_scale, seed=seed, guidance_end=guidance_end ) try: debug_img = Image.open(image.save("/tmp/" + str(seed) + "output.png")) save_image("/tmp/" + str(seed) + "output.png", debug_img) except Exception as e: print("Error 160: " + str(e)) log_params(final_prompt, scale, steps, controlnet_conditioning_scale, guidance_scale, seed, guidance_end, control_image, image) yield f"Completed! Used prompt: {final_prompt}", image, final_prompt except Exception as e: print("Error: " + str(e)) yield f"Error: {str(e)}", None, None with gr.Blocks(title="FLUX Turbo Upscaler", fill_height=True) as demo: gr.Markdown("⚠️ WIP SPACE - UNFINISHED & BUGGY") # status_box = gr.Markdown("🔄 Warming up...") with gr.Row(): with gr.Accordion(): control_image = gr.Image(type="pil", label="Control Image", show_label=False) with gr.Accordion(): generated_image = gr.Image(type="pil", label="Generated Image", format="png", show_label=False) with gr.Row(): with gr.Column(scale=1): prompt = gr.Textbox(lines=4, info="Enter your prompt here or wait for auto-generation...", label="Image Description") focus = gr.Textbox(label="Area(s) of Focus", info="e.g. 'face', 'eyes', 'hair', 'clothes', 'background', etc.", value="clothing material, textures, ethnicity") scale = gr.Slider(1, 3, value=1, label="Scale (Upscale Factor)", step=0.25) with gr.Row(): generate_button = gr.Button("Generate Image", variant="primary") caption_button = gr.Button("Generate Caption", variant="secondary") with gr.Column(scale=1): seed = gr.Slider(0, MAX_SEED, value=42, label="Seed", step=1) steps = gr.Slider(2, 16, value=8, label="Steps", step=1) controlnet_conditioning_scale = gr.Slider(0, 1, value=0.6, label="ControlNet Scale") guidance_scale = gr.Slider(1, 30, value=3.5, label="Guidance Scale") guidance_end = gr.Slider(0, 1, value=1.0, label="Guidance End") with gr.Row(): with gr.Accordion("Auto-Caption settings", open=False, visible=False): system_prompt = gr.Textbox( lines=4, value="Write a straightforward caption for this image. Begin with the main subject and medium. Mention pivotal elements—people, objects, scenery—using confident, definite language. Focus on concrete details like color, shape, texture, and spatial relationships. Show how elements interact. Omit mood and speculative wording. If text is present, quote it exactly. Note any watermarks, signatures, or compression artifacts. Never mention what's absent, resolution, or unobservable details. Vary your sentence structure and keep the description concise, without starting with 'This image is…' or similar phrasing.", label="System Prompt for Captioning", visible=False # Changed to visible ) temperature_slider = gr.Slider( minimum=0.0, maximum=2.0, value=0.6, step=0.05, label="Temperature", info="Higher values make the output more random, lower values make it more deterministic.", visible=False # Changed to visible ) top_p_slider = gr.Slider( minimum=0.0, maximum=1.0, value=0.9, step=0.01, label="Top-p", visible=False # Changed to visible ) max_tokens_slider = gr.Slider( minimum=1, maximum=2048, value=368, step=1, label="Max New Tokens", info="Maximum number of tokens to generate. The model will stop generating if it reaches this limit.", visible=False # Changed to visible ) log_prompt = gr.Checkbox(value=True, label="Log", visible=False) # Changed to visible gr.Markdown("**Tips:** 8 steps is all you need! Incredibly powerful tool, usage instructions coming soon.") caption_state = gr.State() focus_state = gr.State() log_state = gr.State() generate_button.click( fn=process_image, inputs=[ control_image, prompt, system_prompt, scale, steps, controlnet_conditioning_scale, guidance_scale, seed, guidance_end, temperature_slider, top_p_slider, max_tokens_slider, log_prompt ], outputs=[log_state, generated_image, prompt] ) control_image.input( generate_caption, inputs=[control_image], outputs=[caption_state] ).then( generate_focus, inputs=[control_image, focus], outputs=[focus_state] ).then( combine_caption_focus, inputs=[caption_state, focus_state], outputs=[prompt] ) caption_button.click( fn=generate_caption, inputs=[control_image], outputs=[prompt] ).then( generate_focus, inputs=[control_image, focus], outputs=[focus_state] ).then( combine_caption_focus, inputs=[caption_state, focus_state], outputs=[prompt] ) demo.queue().launch(show_error=True)